Fix device of the tensors in set_nms (#10574)

open-mmlab · Jul 21, 2023 · 3edcc16 · 3edcc16
1 parent 00c5586
commit 3edcc16
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/mmdet/models/roi_heads/bbox_heads/multi_instance_bbox_head.py b/mmdet/models/roi_heads/bbox_heads/multi_instance_bbox_head.py
@@ -601,6 +601,10 @@ def set_nms(bboxes: Tensor,
 
         keep = torch.ones(len(ordered_bboxes)) == 1
         ruler = torch.arange(len(ordered_bboxes))
+
+        keep = keep.to(bboxes.device)
+        ruler = ruler.to(bboxes.device)
+
         while ruler.shape[0] > 0:
             basement = ruler[0]
             ruler = ruler[1:]