求教:
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