From 61d6af6705d754f75c54a9c58959bfb2d82a17f5 Mon Sep 17 00:00:00 2001 From: Weijia Liu Date: Tue, 20 Jun 2023 13:50:30 +0800 Subject: [PATCH] [Fix] fix bbox_nms_rotated label device (#873) --- mmrotate/core/post_processing/bbox_nms_rotated.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mmrotate/core/post_processing/bbox_nms_rotated.py b/mmrotate/core/post_processing/bbox_nms_rotated.py index b8abcb93e..4affc2c9c 100644 --- 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)