torch.onnx.export error ------RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument index in method wrapper_CUDA__index_select)
Labels: module: onnx, triaged

/home/xxx/anaconda3/envs/prediction/lib/python3.11/site-packages/torch/onnx/utils.py:689: UserWarning: Constant folding in symbolic shape inference fails: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument index in method wrapper_CUDA__index_select) (Triggered internally at ../torch/csrc/jit/passes/onnx/shape_type_inference.cpp:439.)
  _C._jit_pass_onnx_graph_shape_type_inference(
============= Diagnostic Run torch.onnx.export version 2.0.0+cu118 =============
verbose: False, log level: Level.ERROR
======================= 0 NONE 0 NOTE 0 WARNING 0 ERROR ========================

Traceback (most recent call last):
  File "/home/xxx/depoly_and_infer/convert_onnx.py", line 15857, in <module>
    torch.onnx.export(
  File "/home/xxx/anaconda3/envs/prediction/lib/python3.11/site-packages/torch/onnx/utils.py", line 506, in export
    _export(
  File "/home/xxx/anaconda3/envs/prediction/lib/python3.11/site-packages/torch/onnx/utils.py", line 1548, in _export
    graph, params_dict, torch_out = _model_to_graph(
                                    ^^^^^^^^^^^^^^^^
  File "/home/xxx/anaconda3/envs/prediction/lib/python3.11/site-packages/torch/onnx/utils.py", line 1180, in _model_to_graph
    params_dict = _C._jit_pass_onnx_constant_fold(
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument index in method wrapper_CUDA__index_select)


code here:


   #step 1, ok  
   prob , coord = model(target_history_mcg_input_data.cuda(),
         target_history_lstm_data.cuda(),
         target_history_lstm_data_diff.cuda(),
         other_history_mcg_input_data.cuda(),
         other_agent_history_scatter_numbers.cuda(),
         other_agent_history_scatter_idx.cuda(),
         other_history_lstm_data.cuda(),
         other_history_lstm_data_diff.cuda(),
         road_network_embeddings.cuda(),
         road_network_scatter_numbers.cuda(),
         road_network_scatter_idx.cuda())
   
   print(prob)
   print(coord)   

   #step 2,  error
   torch.onnx.export(
        model,
        (target_history_mcg_input_data.cuda(),
         target_history_lstm_data.cuda(),
         target_history_lstm_data_diff.cuda(),
         other_history_mcg_input_data.cuda(),
         other_agent_history_scatter_numbers.cuda(),
         other_agent_history_scatter_idx.cuda(),
         other_history_lstm_data.cuda(),
         other_history_lstm_data_diff.cuda(),
         road_network_embeddings.cuda(),
         road_network_scatter_numbers.cuda(),
         road_network_scatter_idx.cuda()),
        onnx_path,
        input_names=["target_history_mcg_input_data", 
                     "target_history_lstm_data",
                     "target_history_lstm_data_diff",
                     "other_history_mcg_input_data",
                     "other_agent_history_scatter_numbers",
                     "other_agent_history_scatter_idx",
                     "other_history_lstm_data",
                     "other_history_lstm_data_diff",
                     "road_network_embeddings",
                     "road_network_scatter_numbers",
                     "road_network_scatter_idx"],
        output_names=["probas", "coordinates"],
        dynamic_axes= None,
        opset_version= 11,
        custom_opsets={"torch_scatter":11},
        # verbose=True,
      )