求教:

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