use slice instead of torch.select()

2026-04-26 18:32:35 +08:00 · 2023-04-11 20:26:24 -06:00 · 2023-04-11 20:26:24 -06:00 · 4598d8967a
commit 4598d8967a
parent 9ec7729bea
1 changed files with 1 additions and 1 deletions
--- a/nodes.py
+++ b/nodes.py
@ -1076,7 +1076,7 @@ class ImageToMask:

    def image_to_mask(self, image, channel):
        channels = ["red", "green", "blue"]
-        mask = torch.select(image[0], 2, channels.index(channel))
+        mask = image[0, :, :, channels.index(channel)]
        return (mask,)

 class MaskToImage: