Repository: pytorch
Issue: RuntimeError: rois.device().is_cpu() ASSERT FAILED at /vision/torchvision/csrc/cpu/ROIAlign_cpu.cpp:386
Labels: triaged,module: vision

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->
```
/run/media/xxx/works/python/github/mmdetection/mmdet/core/bbox/transforms.py:161: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if bboxes.size(0) > 0:
/run/media/xxx/works/python/github/mmdetection/mmdet/core/bbox/transforms.py:162: TracerWarning: new_full is a legacy constructor and is not supported in the JIT.
  img_inds = bboxes.new_full((bboxes.size(0), 1), img_id)
/run/media/xxx/works/python/github/mmdetection/mmdet/models/roi_extractors/single_level.py:98: TracerWarning: new_zeros is a legacy constructor and is not supported in the JIT.
  rois.size(0), self.out_channels, *out_size)
/run/media/xxx/works/python/github/mmdetection/mmdet/models/roi_extractors/single_level.py:103: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if inds.any():
Traceback (most recent call last):
  File "/run/media/xxx/works/python/github/mmdetection/tools/pytorch2onnx.py", line 123, in <module>
    main()
  File "/run/media/xxx/works/python/github/mmdetection/tools/pytorch2onnx.py", line 115, in main
    onnx_model = export_onnx_model(model, (input_data, ), args.passes)
  File "/run/media/xxx/works/python/github/mmdetection/tools/pytorch2onnx.py", line 42, in export_onnx_model
    operator_export_type=OperatorExportTypes.ONNX_ATEN_FALLBACK,
  File "/home/xxx/anaconda3/envs/cv-env/lib/python3.6/site-packages/torch/onnx/__init__.py", line 25, in export
    return utils.export(*args, **kwargs)
  File "/home/xxx/anaconda3/envs/cv-env/lib/python3.6/site-packages/torch/onnx/utils.py", line 131, in export
    strip_doc_string=strip_doc_string)
  File "/home/xxx/anaconda3/envs/cv-env/lib/python3.6/site-packages/torch/onnx/utils.py", line 363, in _export
    _retain_param_name, do_constant_folding)
  File "/home/xxx/anaconda3/envs/cv-env/lib/python3.6/site-packages/torch/onnx/utils.py", line 266, in _model_to_graph
    graph, torch_out = _trace_and_get_graph_from_model(model, args, training)
  File "/home/xxx/anaconda3/envs/cv-env/lib/python3.6/site-packages/torch/onnx/utils.py", line 225, in _trace_and_get_graph_from_model
    trace, torch_out = torch.jit.get_trace_graph(model, args, _force_outplace=True)
  File "/home/xxx/anaconda3/envs/cv-env/lib/python3.6/site-packages/torch/jit/__init__.py", line 231, in get_trace_graph
    return LegacyTracedModule(f, _force_outplace, return_inputs)(*args, **kwargs)
  File "/home/xxx/anaconda3/envs/cv-env/lib/python3.6/site-packages/torch/nn/modules/module.py", line 493, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/xxx/anaconda3/envs/cv-env/lib/python3.6/site-packages/torch/jit/__init__.py", line 294, in forward
    out = self.inner(*trace_inputs)
  File "/home/xxx/anaconda3/envs/cv-env/lib/python3.6/site-packages/torch/nn/modules/module.py", line 491, in __call__
    result = self._slow_forward(*input, **kwargs)
  File "/home/xxx/anaconda3/envs/cv-env/lib/python3.6/site-packages/torch/nn/modules/module.py", line 481, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "/run/media/xxx/works/python/github/mmdetection/mmdet/models/detectors/cascade_rcnn.py", line 135, in forward_dummy
    x[:self.bbox_roi_extractor[i].num_inputs], rois)
  File "/home/xxx/anaconda3/envs/cv-env/lib/python3.6/site-packages/torch/nn/modules/module.py", line 491, in __call__
    result = self._slow_forward(*input, **kwargs)
  File "/home/xxx/anaconda3/envs/cv-env/lib/python3.6/site-packages/torch/nn/modules/module.py", line 481, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "/run/media/xxx/works/python/github/mmdetection/mmdet/core/fp16/decorators.py", line 127, in new_func
    return old_func(*args, **kwargs)
  File "/run/media/xxx/works/python/github/mmdetection/mmdet/models/roi_extractors/single_level.py", line 105, in forward
    roi_feats_t = self.roi_layers[i](feats[i], rois_)
  File "/home/xxx/anaconda3/envs/cv-env/lib/python3.6/site-packages/torch/nn/modules/module.py", line 491, in __call__
    result = self._slow_forward(*input, **kwargs)
  File "/home/xxx/anaconda3/envs/cv-env/lib/python3.6/site-packages/torch/nn/modules/module.py", line 481, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "/run/media/xxx/works/python/github/mmdetection/mmdet/ops/roi_align/roi_align.py", line 77, in forward
    self.spatial_scale, self.sample_num)
  File "/home/xxx/anaconda3/envs/cv-env/lib/python3.6/site-packages/torchvision/ops/roi_align.py", line 67, in roi_align
    return _RoIAlignFunction.apply(input, rois, output_size, spatial_scale, sampling_ratio)
  File "/home/xxx/anaconda3/envs/cv-env/lib/python3.6/site-packages/torchvision/ops/roi_align.py", line 23, in forward
    output_size[0], output_size[1], sampling_ratio)
RuntimeError: rois.device().is_cpu() ASSERT FAILED at /vision/torchvision/csrc/cpu/ROIAlign_cpu.cpp:386, please report a bug to PyTorch. rois must be a CPU tensor (ROIAlign_forward_cpu at /vision/torchvision/csrc/cpu/ROIAlign_cpu.cpp:386)
frame #0: std::function<std::string ()>::operator()() const + 0x11 (0x7f9e573bc441 in /home/xxx/anaconda3/envs/cv-env/lib/python3.6/site-packages/torch/lib/libc10.so)
frame #1: c10::Error::Error(c10::SourceLocation, std::string const&) + 0x2a (0x7f9e573bbd7a in /home/xxx/anaconda3/envs/cv-env/lib/python3.6/site-packages/torch/lib/libc10.so)
frame #2: ROIAlign_forward_cpu(at::Tensor const&, at::Tensor const&, float, int, int, int) + 0x269 (0x7f9df8c454d7 in /home/xxx/anaconda3/envs/cv-env/lib/python3.6/site-packages/torchvision/_C.cpython-36m-x86_64-linux-gnu.so)
frame #3: ROIAlign_forward(at::Tensor const&, at::Tensor const&, float, int, int, int) + 0x92 (0x7f9df8c0d60c in /home/xxx/anaconda3/envs/cv-env/lib/python3.6/site-packages/torchvision/_C.cpython-36m-x86_64-linux-gnu.so)
frame #4: <unknown function> + 0x7f75e (0x7f9df8c3375e in /home/xxx/anaconda3/envs/cv-env/lib/python3.6/site-packages/torchvision/_C.cpython-36m-x86_64-linux-gnu.so)
frame #5: <unknown function> + 0x7b4c7 (0x7f9df8c2f4c7 in /home/xxx/anaconda3/envs/cv-env/lib/python3.6/site-packages/torchvision/_C.cpython-36m-x86_64-linux-gnu.so)
frame #6: <unknown function> + 0x76114 (0x7f9df8c2a114 in /home/xxx/anaconda3/envs/cv-env/lib/python3.6/site-packages/torchvision/_C.cpython-36m-x86_64-linux-gnu.so)
frame #7: <unknown function> + 0x763b1 (0x7f9df8c2a3b1 in /home/xxx/anaconda3/envs/cv-env/lib/python3.6/site-packages/torchvision/_C.cpython-36m-x86_64-linux-gnu.so)
frame #8: <unknown function> + 0x65bc4 (0x7f9df8c19bc4 in /home/xxx/anaconda3/envs/cv-env/lib/python3.6/site-packages/torchvision/_C.cpython-36m-x86_64-linux-gnu.so)
<omitting python frames>
frame #15: THPFunction_apply(_object*, _object*) + 0x6b1 (0x7f9e3a384481 in /home/xxx/anaconda3/envs/cv-env/lib/python3.6/site-packages/torch/lib/libtorch_python.so)
```

## To Reproduce

Steps to reproduce the behavior:

1. `python /run/media/xxx/works/python/github/mmdetection/tools/pytorch2onnx.py configs/cascade_rcnn_r50_fpn_1x.py work_dirs/cascade_rcnn_r50_fpn_1xh/latest.pth --out model.onnx`

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->
I am performing torch model to onnx
## Environment
```
PyTorch version: 1.1.0
Is debug build: No
CUDA used to build PyTorch: 9.0.176

OS: Manjaro Linux
GCC version: (GCC) 7.2.1 20180116
CMake version: Could not collect

Python version: 3.6
Is CUDA available: Yes
CUDA runtime version: 10.0.130
GPU models and configuration: GPU 0: GeForce GTX 1050 with Max-Q Design
Nvidia driver version: 440.82
cuDNN version: Could not collect

Versions of relevant libraries:
[pip3] numpy==1.12.1
[pip3] numpydoc==0.6.0
[conda] numpy                     1.18.1                   pypi_0    pypi
[conda] torch                     1.1.0                    pypi_0    pypi
[conda] torchvision               0.3.0                    pypi_0    pypi
```
## Additional context

<!-- Add any other context about the problem here. -->


cc @fmassa