use slice instead of torch.select()

e1d289c1 · missionfloyd · GitHub · e12fb88b · e1d289c1
Unverified Commit e1d289c1 authored Apr 11, 2023 by missionfloyd Committed by GitHub Apr 11, 2023
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

nodes.py nodes.py +1 -1

No files found.
--- 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: