🐛 [Bug] Tensors on CPU/CUDA issue encountered in TorchScript

[gs-olive]

Bug Description

When compiling the EleutherAI GPT-J-6B model using Torch-TensorRT, the following error is encountered:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

To Reproduce

Steps to reproduce the behavior:

1. Run torch_tensorrt.compile with the above model as input, using fp32 precision.
2. Choose a fixed input Tensor of size (1, 358) and data type Int32, and enable truncate_long_and_double with 8 GB workspace.

Expected behavior

Model should successfully compile to Torch-TRT. Specifically, this CPU-tensor error should not arise during compilation.

Environment

• Torch-TensorRT Version (e.g. 1.0.0): bf4474d
• PyTorch Version (e.g. 1.0): 2.1.0.dev20230419+cu118

Additional Context

See #1123 and the fix in #1416 for similar cases.

[github-actions]

This issue has not seen activity for 90 days, Remove stale label or comment or this will be closed in 10 days