🐛 [Bug] Get an Input not on GPU warning when doing the inference

[binliunls]

Bug Description

When run the inference with a converted TensorRT torchscript of MONAI generative model, it reports the warning shown below.

To Reproduce

Steps to reproduce the behavior:

The way to reproduce the error:

1. Build the torch_tensorrt docker from main branch with TENSORRT_VERSION=8.6
2. Start a container with the torch_tensorrt image
3. Clone this branch of MONAI https://github.com/binliunls/MONAI/tree/6838-support-generative-and-hovernet-with-TensorRT
4. Go into the cloned MONAI folder and run `python setup.py develop; pip install -r requirements-dev.txt'
5. Run the python -m monai.bundle download brats_mri_axial_slices_generative_diffusion --bundle_dir ./ to download the model to a local path.
6. Go into the brats_mri_axial_slices_generative_diffusion folder
7. Run the command python -m monai.bundle trt_export --net_id network_def --filepath models/model_trt.ts --ckpt_file models/model.pt --meta_file configs/metadata.json --config_file configs/inference.json --precision fp32 --use_trace "True" --input_shape "[[1, 1, 64, 64], [1,]]" --converter_kwargs "{'truncate_long_and_double': True}" to convert the model

Run the inference with code like:

import torch

...
input_shapes = ([1, 1, 64, 64], [1,])
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = torch.jit.load("/path/to/exported/model_trt.ts")
inputs = [torch.rand(shape, dtype=torch.float32, device=device) for shape in input_shapes]
pred = model(*inputs)
...

Expected behavior

No need to move tensors from CPU to GPU.

Environment

Build information about Torch-TensorRT can be found by turning on debug messages

• TensorRT: 8.6.1+cuda12.0
• Torch-TensorRT Version: 1.4.0
• CPU Architecture: x86-64
• OS: ubuntu 20.04
• Python version:3.8.10
• CUDA version: 12.1
• GPU models and configuration: A100 80G

Additional context

[apbose]

Hi in this case I see that the inputs are getting moved to "cuda" device if it finds the target cuda device. Did you try explicitly placing the tensor to ".cuda()"?

[apbose]

Hi @binliunls , I tried running the above steps to test the Runtime warning with dynamo path, but I run into ModuleNotFoundError: Cannot locate class or function path: 'generative.networks.nets.DiffusionModelUNet'. Is there something missing from requirements.txt?

[binliunls]

Hi @apbose, very appreciate for your response. I think you the libraries shown below should be installed as shown in the metadata.json file.

"monai-generative": "0.2.2"

Thanks, Bin