RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same, Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

Hi Team, when we run the below script to convert the Semantic segmentation Pytorch model to ONNX, We get “RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!” with the latest Pytorch version (V2.0.0) but when I downgrade the Pytorch version to (V1.11.0) the .pth model gets converted to ONNX. Attaching the logs for successful conversion

BUT, when I do that same with Transformation based Panoptic segmentation model, I get same error as mentioned above. I tried downgrading and upgrading Pytorch version but still the issue exists. I also checked by debugging the input and weight tensors where their device shows cuda. Below are the error logs. Kindly assist if this is something related to version issue. Thanks

ERROR1:

frame #38: python3() [0x64ca48]
frame #39: python3() [0x64cb14]
frame #40: python3() [0x64cc0c]
frame #44: __libc_start_main + 0xe8 (0xffff9393ae10 in /lib/aarch64-linux-gnu/libc.so.6)
frame #45: python3() [0x598f48]
(function ComputeConstantFolding)
Traceback (most recent call last):
File “pth_to_onn_full_model.py”, line 103, in
main()
File “pth_to_onn_full_model.py”, line 99, in main
torch.onnx.export(model, dummy_input, export_model_path, verbose=True, input_names=input_names,
File “/usr/local/lib/python3.8/dist-packages/torch/onnx/init.py”, line 305, in export
return utils.export(model, args, f, export_params, verbose, training,
File “/usr/local/lib/python3.8/dist-packages/torch/onnx/utils.py”, line 118, in export
_export(model, args, f, export_params, verbose, training, input_names, output_names,
File “/usr/local/lib/python3.8/dist-packages/torch/onnx/utils.py”, line 719, in _export
_model_to_graph(model, args, verbose, input_names,
File “/usr/local/lib/python3.8/dist-packages/torch/onnx/utils.py”, line 537, in _model_to_graph
params_dict = torch._C._jit_pass_onnx_constant_fold(graph, params_dict,
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument index in method wrapper__index_select)

ERROR 2:
/home/magna1/anaconda3/envs/magp/bin/python /home/magna1/PycharmProjects/Panoptic_Segmentation_Optimization/transformer_pth_to_onnx.py
cuda
Traceback (most recent call last):
File “/home/magna1/PycharmProjects/Panoptic_Segmentation_Optimization/transformer_pth_to_onnx.py”, line 113, in
main()
File “/home/magna1/PycharmProjects/Panoptic_Segmentation_Optimization/transformer_pth_to_onnx.py”, line 84, in main
pixel_level_module_output = model.model.pixel_level_module(pixel_values=dummy_input,
File “/home/magna1/anaconda3/envs/magp/lib/python3.8/site-packages/torch/nn/modules/module.py”, line 1110, in _call_impl
return forward_call(*input, **kwargs)
File “/home/magna1/PycharmProjects/Panoptic_Segmentation_Optimization/lib/models/magna_mask2former.py”, line 20, in forward
last_hidden_feature, backbone_features = self.encoder(pixel_values)
File “/home/magna1/anaconda3/envs/magp/lib/python3.8/site-packages/torch/nn/modules/module.py”, line 1110, in _call_impl
return forward_call(*input, **kwargs)
File “/home/magna1/PycharmProjects/Panoptic_Segmentation_Optimization/lib/models/ti_backbone/multi_input_net.py”, line 139, in forward
x = block(x)
File “/home/magna1/anaconda3/envs/magp/lib/python3.8/site-packages/torch/nn/modules/module.py”, line 1110, in _call_impl
return forward_call(*input, **kwargs)
File “/home/magna1/anaconda3/envs/magp/lib/python3.8/site-packages/torch/nn/modules/container.py”, line 141, in forward
input = module(input)
File “/home/magna1/anaconda3/envs/magp/lib/python3.8/site-packages/torch/nn/modules/module.py”, line 1110, in _call_impl
return forward_call(*input, **kwargs)
File “/home/magna1/anaconda3/envs/magp/lib/python3.8/site-packages/torch/nn/modules/conv.py”, line 447, in forward
return self._conv_forward(input, self.weight, self.bias)
File “/home/magna1/anaconda3/envs/magp/lib/python3.8/site-packages/torch/nn/modules/conv.py”, line 443, in _conv_forward
return F.conv2d(input, weight, bias, self.stride,
RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same

Could you try to script the model before passing it to ONNX as described here?


We tried as u mentioned above!! But we are getting this error in one of our file.
Error:torch.jit.frontend.NotSupportedError: Comprehension ifs are not supported yet

An attachment of the error is added for the reference, Please find it!!
Thank you

Rewrite the tuple comprehension as the error explains if conditions are not supported.

We rewrote the if conditions, then we got errors in “forward” methods.
As mentioned in official site of Pytorch(jit_script),
" nn.Module by default will compile the forward method and recursively compile any methods, submodules, and functions called by forward" so we did @torch.jit.ignore for “forward” methods.

Later since we had try-except in our code we changed these to if-else and some parameter it wanted in type tensor so we defined it as torch.Tensor.

After doing all these changes, now we are getting this error saying OSError: could not get source code
I have added the screenshot, please refer it

Can someone please help and provide any solutions