Fix another instance of the same bug

pangzhangdeng · Apr 10, 2023 · b028d54 · b028d54
1 parent 6325eb8
commit b028d54
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/segment_anything/automatic_mask_generator.py b/segment_anything/automatic_mask_generator.py
@@ -357,7 +357,7 @@ def postprocess_small_regions(
         keep_by_nms = batched_nms(
             boxes.float(),
             torch.as_tensor(scores),
-            torch.zeros(len(boxes)),  # categories
+            torch.zeros_like(boxes[:,0]),  # categories
             iou_threshold=nms_thresh,
         )