From 52215ae06743c00e928905fb8ae5ee2e5fe7bde1 Mon Sep 17 00:00:00 2001
From: Antoine Delplace <34864698+antoinedelplace@users.noreply.github.com>
Date: Wed, 29 Nov 2023 16:30:23 +0100
Subject: [PATCH] Fix permute bug in order to have several output masks and
 images

Fix permute bug in order to have several output masks and images
---
 node.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/node.py b/node.py
index 803fdc5..570e3d2 100644
--- 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)