diff --git a/birefnetNode.py b/birefnetNode.py
index cd5bfaf..8021fdf 100644
--- a/birefnetNode.py
+++ b/birefnetNode.py
@@ -199,11 +199,14 @@ class RembgByBiRefNet:
             image = apply_mask_to_image(image.cpu(), mask.cpu())
 
             _images.append(image)
-            _masks.append(mask)
+            _masks.append(mask.squeeze(0))
 
         out_images = torch.cat(_images, dim=0)
         out_masks = torch.cat(_masks, dim=0)
 
+        if out_masks.shape[0] == 1:
+            out_masks = out_masks.squeeze(0)
+
         return out_images, out_masks