Revert fix for batch size 1

lldacing · Oct 1, 2024 · 09bb745 · 09bb745
1 parent 03ff6c2
commit 09bb745
Showing 1 changed file with 0 additions and 3 deletions.
diff --git a/birefnetNode.py b/birefnetNode.py
@@ -204,9 +204,6 @@ def rem_bg(self, model, images):
         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