🐛 [Bug] Compiling fails with CUDNN_STATUS_NOT_INITIALIZED on Nvidia Orin (aarch64) using NGC Pytorch 24.01 container

Bug Description

To Reproduce

def export_torch_to_trt_16(path_pth, path_trt16_save):
"""
Function exports the pytorch model to tensorrt using torch-tensorrt at FP16
"""

# model = dann_net.DANN_v2(2, 16).cuda()
model = dann_net.DANN_v3(2, 16).cuda()
model.load_state_dict(torch.load(path_pth, map_location="cuda:0"))
model.eval().half()

inp = torchtrt.Input(min_shape=(256, 2, 16),
                    opt_shape=(512, 2, 16),
                    max_shape=(576, 2, 16),
                    dtype=torch.half)

trt_model_fp16 = torchtrt.compile(module=model, inputs=[torchtrt.Input((512, 2, 16), dtype=torch.half)], 
                                  enabled_precisions={torch.half},)
trt_model_fp16.save(path_trt16_save)

Steps to reproduce the behavior:

Expected behavior

Environment

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

• Torch-TensorRT Version (e.g. 1.0.0): 2.2.0a0
• PyTorch Version (e.g. 1.0): 2.2.0a0+81ea7a4
• CPU Architecture: aarch64
• OS (e.g., Linux): Linux
• How you installed PyTorch (conda, pip, libtorch, source): NGC container
• Build command you used (if compiling from source):
• Are you using local sources or building from archives:
• Python version: 3.10
• CUDA version: 12.3
• GPU models and configuration: Nvidia Orin
• Any other relevant information:

Additional context

INFO:torch_tensorrt._compile:ir was set to default, using dynamo as ir
Traceback (most recent call last):
File "/usr/local/lib/python3.10/dist-packages/torch/fx/passes/shape_prop.py", line 151, in run_node
result = super().run_node(n)
File "/usr/local/lib/python3.10/dist-packages/torch/fx/interpreter.py", line 195, in run_node
return getattr(self, n.op)(n.target, args, kwargs)
File "/usr/local/lib/python3.10/dist-packages/torch/fx/interpreter.py", line 267, in call_function
return target(*args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/torch/_ops.py", line 513, in call
return self._op(*args, **kwargs or {})
RuntimeError: cuDNN error: CUDNN_STATUS_NOT_INITIALIZED
WARNING:torch_tensorrt.dynamo.lowering.passes.fuse_prims_broadcast:Shape Propagation Failed on Graph, skipping fuse_prims_broadcast lowering pass
Traceback (most recent call last):
File "/usr/local/lib/python3.10/dist-packages/torch/fx/passes/shape_prop.py", line 151, in run_node
result = super().run_node(n)
File "/usr/local/lib/python3.10/dist-packages/torch/fx/interpreter.py", line 195, in run_node
return getattr(self, n.op)(n.target, args, kwargs)
File "/usr/local/lib/python3.10/dist-packages/torch/fx/interpreter.py", line 267, in call_function
return target(*args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/torch/_ops.py", line 513, in call
return self._op(*args, **kwargs or {})
RuntimeError: cuDNN error: CUDNN_STATUS_NOT_INITIALIZED

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "/usr/local/lib/python3.10/dist-packages/torch_tensorrt/dynamo/lowering/passes/fuse_prims_broadcast.py", line 22, in fuse_prims_broadcast
ShapeProp(gm).propagate(*sample_inputs)
File "/usr/local/lib/python3.10/dist-packages/torch/fx/passes/shape_prop.py", line 193, in propagate
return super().run(*fake_args)
File "/usr/local/lib/python3.10/dist-packages/torch/fx/interpreter.py", line 138, in run
self.env[node] = self.run_node(node)
File "/usr/local/lib/python3.10/dist-packages/torch/fx/passes/shape_prop.py", line 156, in run_node
raise RuntimeError(
RuntimeError: ShapeProp error for: node=%convolution : [num_users=1] = call_function[target=torch.ops.aten.convolution.default](args = (%arg25_1, %res1_c1_weight, %res1_c1_bias, [1], [2], [1], False, [0], 1), kwargs = {}) with meta={'stack_trace': ' File "/home/vrs/Documents/doae/root2/neural_networks/dann_net.py", line 280, in forward\n x = self.res1(x_phases)\n File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1519, in call_impl\n return forward_call(*args, **kwargs)\n File "/home/vrs/Documents/doae/root2/neural_networks/blocks.py", line 55, in forward\n x_out = self.bn1(self.act1(self.c1(x)))\n', 'nn_module_stack': {'L__self_': ('', <class 'neural_networks.dann_net.DANN_v3'>), 'L__self___res1': ('res1', <class 'neural_networks.blocks.Conv1DCatResBlockNoPool'>), 'L__self___res1_c1': ('res1.c1', <class 'torch.nn.modules.conv.Conv1d'>)}, 'source_fn_stack': [('l__self___res1_c1', <class 'torch.nn.modules.conv.Conv1d'>)], 'original_aten': <OpOverload(op='aten.convolution', overload='default')>, 'from_node': [('l__self___res1_c1', 'L__self___res1_c1'), ('convolution', <OpOverload(op='aten.convolution', overload='default')>)], 'seq_nr': 93, 'val': FakeTensor(..., device='cuda:0', size=(512, 32, 16), dtype=torch.float16), 'tensor_meta': TensorMetadata(shape=torch.Size([512, 32, 16]), dtype=torch.float16, requires_grad=False, stride=(512, 16, 1), memory_format=torch.contiguous_format, is_quantized=False, qparams={})}

While executing %convolution : [num_users=1] = call_function[target=torch.ops.aten.convolution.default](args = (%arg25_1, %res1_c1_weight, %res1_c1_bias, [1], [2], [1], False, [0], 1), kwargs = {})
Original traceback:
File "/home/vrs/Documents/doae/root2/neural_networks/dann_net.py", line 280, in forward
x = self.res1(x_phases)
File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1519, in _call_impl
return forward_call(*args, **kwargs)
File "/home/vrs/Documents/doae/root2/neural_networks/blocks.py", line 55, in forward
x_out = self.bn1(self.act1(self.c1(x)))

INFO:torch_tensorrt.dynamo._compiler:Compilation Settings: CompilationSettings(precision=torch.float16, debug=False, workspace_size=0, min_block_size=5, torch_executed_ops=[], pass_through_build_failures=False, max_aux_streams=None, version_compatible=False, optimization_level=None, use_python_runtime=False, truncate_long_and_double=False, use_fast_partitioner=True, enable_experimental_decompositions=False, device=Device(type=DeviceType.GPU, gpu_id=0), require_full_compilation=False, disable_tf32=False, sparse_weights=False, refit=False, engine_capability=<EngineCapability.DEFAULT: 0>, num_avg_timing_iters=1, dla_sram_size=1048576, dla_local_dram_size=1073741824, dla_global_dram_size=536870912)

Traceback (most recent call last):
File "/home/vrs/Documents/doae/root2/deploy.py", line 16, in
exporters.export_torch_to_trt_16(path_pth="/home/vrs/Documents/doae/root2/trained_models/pth/DANN_v3_27_6_2024_1490.pth",
File "/home/vrs/Documents/doae/root2/tools/exporters.py", line 31, in export_torch_to_trt_16
trt_model_fp16 = torchtrt.compile(module=model, inputs=[torchtrt.Input((512, 2, 16), dtype=torch.half)],
File "/usr/local/lib/python3.10/dist-packages/torch_tensorrt/_compile.py", line 228, in compile
trt_graph_module = dynamo_compile(
File "/usr/local/lib/python3.10/dist-packages/torch_tensorrt/dynamo/_compiler.py", line 222, in compile
return compile_module(gm, inputs, settings)
File "/usr/local/lib/python3.10/dist-packages/torch_tensorrt/dynamo/_compiler.py", line 301, in compile_module
submodule_inputs = partitioning.get_submod_inputs(
File "/usr/local/lib/python3.10/dist-packages/torch_tensorrt/dynamo/partitioning/common.py", line 84, in get_submod_inputs
mod(*torch_inputs)
File "/usr/local/lib/python3.10/dist-packages/torch/fx/graph_module.py", line 736, in call_wrapped
return self._wrapped_call(self, *args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/torch/fx/graph_module.py", line 315, in call
raise e
File "/usr/local/lib/python3.10/dist-packages/torch/fx/graph_module.py", line 302, in call
return super(self.cls, obj).call(*args, **kwargs) # type: ignore[misc]
File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1510, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1519, in _call_impl
return forward_call(*args, **kwargs)
File "<eval_with_key>.28", line 5, in forward
File "/usr/local/lib/python3.10/dist-packages/torch/fx/graph_module.py", line 736, in call_wrapped
return self._wrapped_call(self, *args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/torch/fx/graph_module.py", line 315, in call
raise e
File "/usr/local/lib/python3.10/dist-packages/torch/fx/graph_module.py", line 302, in call
return super(self.cls, obj).call(*args, **kwargs) # type: ignore[misc]
File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1510, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1560, in _call_impl
result = forward_call(*args, **kwargs)
File "<eval_with_key>.27", line 7, in forward
File "/usr/local/lib/python3.10/dist-packages/torch/_ops.py", line 513, in call
return self._op(*args, **kwargs or {})
RuntimeError: cuDNN error: CUDNN_STATUS_NOT_INITIALIZED

Start with tools/exporters.py at export_torch_to_trt_16 and deploy.py at the failing torchtrt.compile call. Review the model path through neural_networks/dann_net.py and neural_networks/blocks.py, then reproduce the CUDNN_STATUS_NOT_INITIALIZED failure in the stated NGC container on Nvidia Orin. Done means the reported compilation failure is explained and a verified resolution or clear compatibility finding is documented.