Add a batch dimension to masks loaded by LoadMaskBase64

2024-06-03 23:16:59 +02:00
parent 7c338bbd2a
commit afefeca48f
1 changed files with 1 additions and 1 deletions
@@ -49,7 +49,7 @@ class LoadMaskBase64:
        img = torch.from_numpy(img)
        if img.dim() == 3:  # RGB(A) input, use red channel
            img = img[:, :, 0]
-        return (img,)
+        return (img.unsqueeze(0),)


 class SendImageWebSocket: