From b028d54358c6d4b80b1eef27ec98968150c65943 Mon Sep 17 00:00:00 2001
From: Louis Maddox <louismmx@gmail.com>
Date: Thu, 6 Apr 2023 11:12:04 +0100
Subject: [PATCH] Fix another instance of the same bug

---
 segment_anything/automatic_mask_generator.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/segment_anything/automatic_mask_generator.py b/segment_anything/automatic_mask_generator.py
index 9c04d76..d0f5a1b 100644
--- a/segment_anything/automatic_mask_generator.py
+++ b/segment_anything/automatic_mask_generator.py
@@ -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,
         )