[Bug] bbox_nms_rotated.py bugged

[Kiochio]

Prerequisite

• I have searched Issues and Discussions but cannot get the expected help.
• I have read the FAQ documentation but cannot get the expected help.
• The bug has not been fixed in the latest version (master) or latest version (1.x).

Task

I'm using the official example scripts/configs for the officially supported tasks/models/datasets.

Branch

master branch https://github.com/open-mmlab/mmrotate

Environment

sys.platform: win32
Python: 3.11.8 | packaged by Anaconda, Inc. | (main, Feb 26 2024, 21:34:05) [MSC v.1916 64 bit (AMD64)]
CUDA available: True
GPU 0: NVIDIA GeForce RTX 2060 SUPER
CUDA_HOME: C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.8
NVCC: Cuda compilation tools, release 11.8, V11.8.89
MSVC: Compilateur d'optimisation Microsoft (R) C/C++ version 19.35.32217.1 pour x64
GCC: n/a
PyTorch: 2.0.0+cu118
PyTorch compiling details: PyTorch built with:

• C++ Version: 199711
• MSVC 193431937
• Intel(R) Math Kernel Library Version 2020.0.2 Product Build 20200624 for Intel(R) 64 architecture applications
• Intel(R) MKL-DNN v2.7.3 (Git Hash 6dbeffbae1f23cbbeae17adb7b5b13f1f37c080e)
• OpenMP 2019
• LAPACK is enabled (usually provided by MKL)
• CPU capability usage: AVX2
• CUDA Runtime 11.8
• NVCC architecture flags: -gencode;arch=compute_37,code=sm_37;-gencode;arch=compute_50,code=sm_50;-gencode;arch=compute_60,code=sm_60;-gencode;arch=compute_61,code=sm_61;-gencode;arch=compute_70,code=sm_70;-gencode;arch=compute_75,code=sm_75;-gencode;arch=compute_80,code=sm_80;-gencode;arch=compute_86,code=sm_86;-gencode;arch=compute_90,code=sm_90;-gencode;arch=compute_37,code=compute_37
• CuDNN 8.7
• Magma 2.5.4
• Build settings: BLAS_INFO=mkl, BUILD_TYPE=Release, CUDA_VERSION=11.8, CUDNN_VERSION=8.7.0, CXX_COMPILER=C:/actions-runner/_work/pytorch/pytorch/builder/windows/tmp_bin/sccache-cl.exe, CXX_FLAGS=/DWIN32 /D_WINDOWS /GR /EHsc /w /bigobj /FS -DUSE_PTHREADPOOL -DNDEBUG -DUSE_KINETO -DLIBKINETO_NOCUPTI -DLIBKINETO_NOROCTRACER -DUSE_FBGEMM -DUSE_XNNPACK -DSYMBOLICATE_MOBILE_DEBUG_HANDLE, LAPACK_INFO=mkl, PERF_WITH_AVX=1, PERF_WITH_AVX2=1, PERF_WITH_AVX512=1, TORCH_DISABLE_GPU_ASSERTS=OFF, TORCH_VERSION=2.0.0, USE_CUDA=ON, USE_CUDNN=ON, USE_EXCEPTION_PTR=1, USE_GFLAGS=OFF, USE_GLOG=OFF, USE_MKL=ON, USE_MKLDNN=ON, USE_MPI=OFF, USE_NCCL=OFF, USE_NNPACK=OFF, USE_OPENMP=ON, USE_ROCM=OFF,

TorchVision: 0.15.1+cu118
OpenCV: 4.9.0
MMCV: 1.7.2
MMCV Compiler: MSVC 192930148
MMCV CUDA Compiler: 11.8
MMRotate: 0.3.4+e384050

Reproduces the problem - code sample

nothing

Reproduces the problem - command or script

python tools/test.py ${CONFIG_FILE} ${CHECKPOINT_FILE} --eval mAP

Reproduces the problem - error message

File "C:\Users\Kioch\anaconda3\envs\ptcgar\Lib\site-packages\mmrotate\core\post_processing\bbox_nms_rotated.py", line 62, in multiclass_nms_rotated
bboxes, scores, labels = bboxes[inds], scores[inds], labels[inds]
~~~~^^^
RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

Additional information

Is it a normal behavior for my setup? Am I doing something wrong?
I already find the problem. For me the labels variable created line 42 in the file mmrotate/mmrotate/core/post_processing
/bbox_nms_rotated.py is created always on the CPU and in the case where we want to use the GPU for our model the code will crash because at line 58 it tries to use variables that are not on the same device.

A fix that I found is to replace line 42 by:
if bboxes.get_device() == -1: device = 'cpu' else: device = bboxes.get_device() labels = torch.arange(num_classes, dtype=torch.long).to(device)