[Fix] fix bbox_nms_rotated label device (#873)

open-mmlab · Jun 20, 2023 · 61d6af6 · 61d6af6
1 parent e1d96ec
commit 61d6af6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/mmrotate/core/post_processing/bbox_nms_rotated.py b/mmrotate/core/post_processing/bbox_nms_rotated.py
@@ -39,7 +39,7 @@ def multiclass_nms_rotated(multi_bboxes,
             multi_scores.size(0), num_classes, 5)
     scores = multi_scores[:, :-1]
 
-    labels = torch.arange(num_classes, dtype=torch.long)
+    labels = torch.arange(num_classes, dtype=torch.long, device=scores.device)
     labels = labels.view(1, -1).expand_as(scores)
     bboxes = bboxes.reshape(-1, 5)
     scores = scores.reshape(-1)