PyTorch model can forward but onnx export tell me RuntimeError: Expected object of device type cuda but got device type cpu for argument #3 'index' in call to _th_index_select

1

The error is:

Traceback (most recent call last):
  File "export_onnx.py", line 61, in <module>
    output_names=['output'], enable_onnx_checker=False)  # output_names=['output']
  File "/usr/local/lib/python3.6/dist-packages/torch/onnx/__init__.py", line 208, in export
    custom_opsets, enable_onnx_checker, use_external_data_format)
  File "/usr/local/lib/python3.6/dist-packages/torch/onnx/utils.py", line 92, in export
    use_external_data_format=use_external_data_format)
  File "/usr/local/lib/python3.6/dist-packages/torch/onnx/utils.py", line 530, in _export
    fixed_batch_size=fixed_batch_size)
  File "/usr/local/lib/python3.6/dist-packages/torch/onnx/utils.py", line 409, in _model_to_graph
    _export_onnx_opset_version)
RuntimeError: Expected object of device type cuda but got device type cpu for argument #3 'index' in call to _th_index_select

I am tracing a model to onnx, the question is the model can foward in GPU mode, but when trace it tell me that ONNX find a node in CPU, how could that possible?

I mean if there were one node in CPU then how could my model forward???

Any body knows why?