Deploy model to TensorRT failed in conversion to ONNX

[yossicohen-96]

Hello,

I've been attempting to deploy the mask2former_flash_internimage_s_640_160k_ade20k_ss model using the provided deploy.py script from the internimage repository, located in the segmentation folder. I've configured a deployment Python file as follows to integrate TensorRT with ONNX:

_base_ = ['./segmentation_static.py', '../_base_/backends/tensorrt.py']

onnx_config = dict(input_shape=[640, 640])
backend_config = dict(
    common_config=dict(max_workspace_size=1 << 30),
    model_inputs=[
        dict(
            input_shapes=dict(
                input=dict(
                    min_shape=[1, 3, 640, 640],
                    opt_shape=[1, 3, 640, 640],
                    max_shape=[1, 3, 640, 640]
                )
            )
        ]
    )

Additionally, I utilized the corresponding configuration file located at segmentation/configs/ade20k/mask2former_flash_internimage_s_640_160k_ade20k_ss.py.

Despite setting the device to CUDA (--device=cuda), the script fails during the conversion from a Torch graph to an ONNX graph with the following error indicating that the model contains both CPU and GPU tensors:

Traceback (most recent call last):
...
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)
2024-06-10 15:04:59,856 - mmdeploy - ERROR - `mmdeploy.apis.pytorch2onnx.torch2onnx` with Call id: 0 failed. exit.

Could you please help me understand why this issue is occurring and advise on how to correctly deploy the model ensuring all tensors are on the same device?

Thanks

[ChenJian7578]

Have you solved this problem? The problem seems to be because the model and input are one on the cpu and one on the gpu