Fix permute bug in order to have several output masks and images

codeinvain · Nov 29, 2023 · 52215ae · 52215ae
1 parent 906a51f
commit 52215ae
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/node.py b/node.py
@@ -231,7 +231,7 @@ def sam_segment(
         point_labels=None,
         boxes=transformed_boxes.to(sam_device),
         multimask_output=False)
-    masks = masks.permute(1, 0, 2, 3).cpu().numpy()
+    masks = masks.cpu().numpy()
     return create_tensor_output(image_np, masks, boxes)