Fix another instance of the same bug

2023-04-10 11:58:44 -07:00
parent 6325eb8512
commit b028d54358
1 changed files with 1 additions and 1 deletions
@@ -357,7 +357,7 @@ class SamAutomaticMaskGenerator:
        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,
        )