conversion failed with RotatedRetinanet + onnxruntime-gpu

[RunningLeon]

Describe the bug

conversion failed with RotatedRetinanet + onnxruntime-gpu

Reproduction

1. What command or script did you run?

python tools/regression_test.py --codebase mmrotate --backend onnxruntime --models RotatedRetinanet --device cuda:0

2. Did you make any modifications on the code or config? Did you understand what you have modified?

Environment

1. Please run python tools/check_env.py to collect necessary environment information and paste it here.
2. You may add addition that may be helpful for locating the problem, such as
   • How you installed PyTorch [e.g., pip, conda, source]
   • Other environment variables that may be related (such as $PATH, $LD_LIBRARY_PATH, $PYTHONPATH, etc.)

Error traceback

Traceback (most recent call last):
  File "/home/xxx/anaconda3/envs/pt1.8/lib/python3.7/multiprocessing/process.py", line 297, in _bootstrap
    self.run()
  File "/home/xxx/anaconda3/envs/pt1.8/lib/python3.7/multiprocessing/process.py", line 99, in run
    self._target(*self._args, **self._kwargs)
  File "/home/xxx/projects/openmmlab/mmdeploy/mmdeploy/apis/core/pipeline_manager.py", line 107, in __call__
    ret = func(*args, **kwargs)
  File "/home/xxx/projects/openmmlab/mmdeploy/mmdeploy/apis/pytorch2onnx.py", line 105, in torch2onnx
    optimize=optimize)
  File "/home/xxx/projects/openmmlab/mmdeploy/mmdeploy/apis/core/pipeline_manager.py", line 356, in _wrap
    return self.call_function(func_name_, *args, **kwargs)
  File "/home/xxx/projects/openmmlab/mmdeploy/mmdeploy/apis/core/pipeline_manager.py", line 326, in call_function
    return self.call_function_local(func_name, *args, **kwargs)
  File "/home/xxx/projects/openmmlab/mmdeploy/mmdeploy/apis/core/pipeline_manager.py", line 275, in call_function_local
    return pipe_caller(*args, **kwargs)
  File "/home/xxx/projects/openmmlab/mmdeploy/mmdeploy/apis/core/pipeline_manager.py", line 107, in __call__
    ret = func(*args, **kwargs)
  File "/home/xxx/projects/openmmlab/mmdeploy/mmdeploy/apis/onnx/export.py", line 132, in export
    verbose=verbose)
  File "/home/xxx/anaconda3/envs/pt1.8/lib/python3.7/site-packages/torch/onnx/__init__.py", line 276, in export
    custom_opsets, enable_onnx_checker, use_external_data_format)
  File "/home/xxx/anaconda3/envs/pt1.8/lib/python3.7/site-packages/torch/onnx/utils.py", line 94, in export
    use_external_data_format=use_external_data_format)
  File "/home/xxx/anaconda3/envs/pt1.8/lib/python3.7/site-packages/torch/onnx/utils.py", line 698, in _export
    dynamic_axes=dynamic_axes)
  File "/home/xxx/projects/openmmlab/mmdeploy/mmdeploy/core/rewriters/rewriter_utils.py", line 379, in wrapper
    return self.func(self, *args, **kwargs)
  File "/home/xxx/projects/openmmlab/mmdeploy/mmdeploy/apis/onnx/optimizer.py", line 10, in model_to_graph__custom_optimizer
    graph, params_dict, torch_out = ctx.origin_func(*args, **kwargs)
  File "/home/xxx/anaconda3/envs/pt1.8/lib/python3.7/site-packages/torch/onnx/utils.py", line 456, in _model_to_graph
    use_new_jit_passes)
  File "/home/xxx/anaconda3/envs/pt1.8/lib/python3.7/site-packages/torch/onnx/utils.py", line 417, in _create_jit_graph
    graph, torch_out = _trace_and_get_graph_from_model(model, args)
  File "/home/xxx/anaconda3/envs/pt1.8/lib/python3.7/site-packages/torch/onnx/utils.py", line 377, in _trace_and_get_graph_from_model
    torch.jit._get_trace_graph(model, args, strict=False, _force_outplace=False, _return_inputs_states=True)
  File "/home/xxx/anaconda3/envs/pt1.8/lib/python3.7/site-packages/torch/jit/_trace.py", line 1139, in _get_trace_graph
    outs = ONNXTracedModule(f, strict, _force_outplace, return_inputs, _return_inputs_states)(*args, **kwargs)
  File "/home/xxx/anaconda3/envs/pt1.8/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/xxx/anaconda3/envs/pt1.8/lib/python3.7/site-packages/torch/jit/_trace.py", line 130, in forward
    self._force_outplace,
  File "/home/xxx/anaconda3/envs/pt1.8/lib/python3.7/site-packages/torch/jit/_trace.py", line 116, in wrapper
    outs.append(self.inner(*trace_inputs))
  File "/home/xxx/anaconda3/envs/pt1.8/lib/python3.7/site-packages/torch/nn/modules/module.py", line 887, in _call_impl
    result = self._slow_forward(*input, **kwargs)
  File "/home/xxx/anaconda3/envs/pt1.8/lib/python3.7/site-packages/torch/nn/modules/module.py", line 860, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "/home/xxx/projects/openmmlab/mmdeploy/mmdeploy/core/rewriters/rewriter_utils.py", line 379, in wrapper
    return self.func(self, *args, **kwargs)
  File "/home/xxx/projects/openmmlab/mmdeploy/mmdeploy/codebase/mmdet/models/detectors/base.py", line 70, in base_detector__forward
    return __forward_impl(ctx, self, img, img_metas=img_metas, **kwargs)
  File "/home/xxx/projects/openmmlab/mmdeploy/mmdeploy/core/optimizers/function_marker.py", line 261, in g
    rets = f(*args, **kwargs)
  File "/home/xxx/projects/openmmlab/mmdeploy/mmdeploy/codebase/mmdet/models/detectors/base.py", line 26, in __forward_impl
    return self.simple_test(img, img_metas, **kwargs)
  File "/home/xxx/projects/openmmlab/mmdeploy/mmdeploy/core/rewriters/rewriter_utils.py", line 379, in wrapper
    return self.func(self, *args, **kwargs)
  File "/home/xxx/projects/openmmlab/mmdeploy/mmdeploy/codebase/mmrotate/models/single_stage_rotated_detector.py", line 32, in single_stage_rotated_detector__simple_test
    outs = self.bbox_head.get_bboxes(*outs, img_metas, rescale=rescale)
  File "/home/xxx/projects/openmmlab/mmdeploy/mmdeploy/core/rewriters/rewriter_utils.py", line 379, in wrapper
    return self.func(self, *args, **kwargs)
  File "/home/xxx/projects/openmmlab/mmdeploy/mmdeploy/codebase/mmrotate/models/dense_heads/rotated_anchor_head.py", line 138, in rotated_anchor_head__get_bbox
    keep_top_k=keep_top_k)
  File "/home/xxx/projects/openmmlab/mmdeploy/mmdeploy/core/optimizers/function_marker.py", line 261, in g
    rets = f(*args, **kwargs)
  File "/home/xxx/projects/openmmlab/mmdeploy/mmdeploy/codebase/mmrotate/core/post_processing/bbox_nms.py", line 179, in multiclass_nms_rotated
    _multiclass_nms_rotated(*args, **kwargs)
  File "/home/xxx/projects/openmmlab/mmdeploy/mmdeploy/codebase/mmrotate/core/post_processing/bbox_nms.py", line 119, in _multiclass_nms_rotated
    score_threshold)
  File "/home/xxx/projects/openmmlab/mmdeploy/mmdeploy/mmcv/ops/nms_rotated.py", line 48, in forward
    0)
RuntimeError: CUDA error: invalid configuration argument
2022-09-01 15:24:28,473 - mmdeploy - ERROR - `mmdeploy.apis.pytorch2onnx.torch2onnx` with Call id: 0 failed. exit.

Bug fix

If you have already identified the reason, you can provide the information here. If you are willing to create a PR to fix it, please also leave a comment here and that would be much appreciated!