Merge pull request #4 from fAIseh00d/main

Fix output mask redundant dims
2024-10-02 00:56:32 +08:00
parent e3a7de4b8b 09bb7452d9
commit 9108de9551
1 changed files with 1 additions and 1 deletions
@@ -199,7 +199,7 @@ 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)