Fails intermittently when intermediate layers produce large tensors

[naiveen]

Hi, I am facing issues compiling one of my models.
I have a model in which I pass output from intermediate layers through CONVTRANSPOSE1D layers and sample up the input tensor. As a result, the output from these layers quickly grows in size and produces tensors of size (1, 256, 10000).

In such a case,

I run into RuntimeError (More common)

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/naveen/miniconda3/envs/trt/lib/python3.10/site-packages/torch_tensorrt/_compile.py", line 185, in compile
    compiled_ts_module: torch.jit.ScriptModule = torchscript_compile(
  File "/home/naveen/miniconda3/envs/trt/lib/python3.10/site-packages/torch_tensorrt/ts/_compiler.py", line 151, in compile
    compiled_cpp_mod = _C.compile_graph(module._c, _parse_compile_spec(spec))
RuntimeError: [Error thrown at core/conversion/var/Var.cpp:127] Expected isIValue() to be true but got false
Requested IValue from Var, however Var type is nvinfer1::ITensor

A segmentation error and my kernel crashes. (Less common)

INFO:torch_tensorrt._compile:Module was provided as a torch.nn.Module, trying to script the module with torch.jit.script. In the event of a failure please preconvert your module to TorchScript
ERROR: [Torch-TensorRT TorchScript Conversion Context] - 3: [network.cpp::addConvolutionNd::1061] Error Code 3: API Usage Error (Parameter check failed at: optimizer/api/network.cpp::addConvolutionNd::1061, condition: nbSpatialDims == 2 || nbSpatialDims == 3
)
/home/naveen/miniconda3/envs/trt/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 1 leaked semaphore objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
Segmentation fault (core dumped)

Very rarely the compilation is successful.

I tried creating simple networks that generate large intermediate tensors but couldn't replicate the behavior. But, when I change the input to my model in order to contain the outputs from intermediate layer to be under 2000, the compilation is successful with occasional segmentation faults and the runtimeerror is completely avoided.

My env: pytorch 2.2.1+cu121, tensorrt 8.6.1, torch_tensorrt 2.2.0 , RTX 4090
Compilation command : trt_module = torchtrt.compile(module, inputs=[sample_input], workspace_size=1<<30, ir='torchscript')

I encounter the same issues when I run on the latest NGC container. However, when I tried to compile on colab, failures happen less frequently (1 Runtime Error of 6 tests). Please share any thoughts that can help me better debug this issue or what if i am doing anything wrong.

[naiveen]

In case any one facing this issue, having module on cpu worked in compiling my model to tensorrt.