From e1d289c1ec6894e15af0b57b6630b853341c61fa Mon Sep 17 00:00:00 2001 From: missionfloyd Date: Tue, 11 Apr 2023 20:26:24 -0600 Subject: [PATCH] use slice instead of torch.select() --- nodes.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nodes.py b/nodes.py index ecd931d69..815631f58 100644 --- 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: