🐛 [Bug] Error encountered when compiling TFT model

[gs-olive]

Bug Description

When compiling the TFT model with Torch-TRT, the following error is encountered:

File "/usr/local/lib/python3.8/dist-packages/torch_tensorrt/ts/_compiler.py", line 136, in compile
        compiled_cpp_mod = _C.compile_graph(module._c, _parse_compile_spec(spec))
    RuntimeError: [Error thrown at core/util/trt_util.cpp:165] pos >= 0 && pos <= d.nbDims ASSERT FAILED at core/util/trt_util.cpp:165, consider filing a bug: https://www.github.com/NVIDIA/Torch-TensorRT/issues
    ERROR: Index to unsqueeze is out of bounds.

To Reproduce

Steps to reproduce the behavior:

1. Run torch_tensorrt.compile with TFT model as input, using fp32 precision.
2. Choose fixed input sizes and enable truncate_long_and_double with 12 GB workspace.

Expected behavior

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

Environment

• Torch-TensorRT Version (e.g. 1.0.0): 97209fe
• PyTorch Version (e.g. 1.0): 2.0.0.dev20230219+cu117