Move batched NMS indices to correct device (closes #17)

2023-04-06 11:07:51 +01:00 · 2023-04-06 11:07:51 +01:00 · 6325eb8512
parent 2780a301de
commit 6325eb8512
1 changed files with 2 additions and 2 deletions
--- a/segment_anything/automatic_mask_generator.py
+++ b/segment_anything/automatic_mask_generator.py
@ -214,7 +214,7 @@ class SamAutomaticMaskGenerator:
            keep_by_nms = batched_nms(
                data["boxes"].float(),
                scores,
-                torch.zeros(len(data["boxes"])),  # categories
+                torch.zeros_like(data["boxes"][:,0]),  # categories
                iou_threshold=self.crop_nms_thresh,
            )
            data.filter(keep_by_nms)
@ -251,7 +251,7 @@ class SamAutomaticMaskGenerator:
        keep_by_nms = batched_nms(
            data["boxes"].float(),
            data["iou_preds"],
-            torch.zeros(len(data["boxes"])),  # categories
+            torch.zeros_like(data["boxes"][:,0]),  # categories
            iou_threshold=self.box_nms_thresh,
        )
        data.filter(keep_by_nms)