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

lmmx · Apr 6, 2023 · 006dbe5 · 006dbe5
1 parent aac76a1
commit 006dbe5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/segment_anything/automatic_mask_generator.py b/segment_anything/automatic_mask_generator.py
@@ -214,7 +214,7 @@ def _generate_masks(self, image: np.ndarray) -> MaskData:
             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 @@ def _process_crop(
         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)