Cannot export MAISI to TensorRT with dynamic shape

[binliunls]

Describe the bug
Cannot export the controlnet of MAISI to TensorRT given a dynamic shape input, which is common for the controlnet use case.

To Reproduce
Run the script under MONAI 1.4.1rc1 version

from monai.apps.generation.maisi.networks.controlnet_maisi import ControlNetMaisi
from monai.networks.trt_compiler import trt_compile
import torch
def test():
    d_kwargs = {
        "spatial_dims": 3,
        "in_channels": 4,
        "out_channels": 4,
        "num_channels": [
            64,
            128,
            256,
            512
        ],
        "attention_levels": [
            False,
            False,
            True,
            True
        ],
        "num_head_channels": [
            0,
            0,
            32,
            32
        ],
        "num_res_blocks": 2,
        "use_flash_attention": True,
        "include_top_region_index_input": False,
        "include_bottom_region_index_input": False,
        "include_spacing_input": True,
        "num_class_embeds": 128,
        "resblock_updown": True,
        "include_fc": True
    }
    c_kwargs = {"spatial_dims": 3,
        "in_channels": 4,
        "num_channels": [
            64,
            128,
            256,
            512
        ],
        "attention_levels": [
            False,
            False,
            True,
            True
        ],
        "num_head_channels": [
            0,
            0,
            32,
            32
        ],
        "num_res_blocks": 2,
        "use_flash_attention": True,
        "conditioning_embedding_in_channels": 8,
        "conditioning_embedding_num_channels": [
            8,
            32,
            64
        ],
        "num_class_embeds": 128,
        "resblock_updown": True,
        "include_fc": True
        }
    controlnet = ControlNetMaisi(**c_kwargs)
    controlnet.eval()
    controlnet.to("cuda")

    c_trt_args = {
        "export_args": {
            "dynamo": False,
            "report": True,
        },
        "output_lists": [
            [
                -1
            ],
            []
        ],
        "input_profiles": [
            {"x": [[1, 4, 64, 64, 32], [1, 4, 128, 128, 128], [1, 4, 128, 128, 128]]}
        ]
    }
    timesteps = torch.randint(0, 1000, (1,), device=torch.device("cuda")).long()
    controlnet_cond = torch.rand((1, 8, 256, 256, 256), device=torch.device("cuda"))
    context = torch.randn((1, 1, 2))
    trt_control_net = trt_compile(controlnet, "/home/liubin/tmp", c_trt_args)
    input_noise = torch.rand((1, 4, 64, 64, 64), device=torch.device("cuda"))
    output = trt_control_net(input_noise, timesteps=timesteps, controlnet_cond=controlnet_cond, class_labels=torch.randint(1,2,(1,),device=torch.device("cuda")))
    print(output.shape())

test()

Expected behavior
Can export the controlnet to TensorRT with a dynamic input shape.

Error report

2025-02-18 13:48:00,722 - INFO - Loading TensorRT engine: /tmp.plan
[W] 'colored' module is not installed, will not use colors when logging. To enable colors, please install the 'colored' module: python3 -m pip install colored
[I] Loading bytes from /tmp.plan
2025-02-18 13:48:00,722 - INFO - Exception while loading the engine:
[Errno 2] No such file or directory: '/tmp.plan'
Adding casts around norms...
Swapped 0 modules
2025-02-18 13:48:00,723 - INFO - Exporting to /tmp/tmpza391yxe/model.onnx:
unrolled_inputs=['x', 'timesteps', 'controlnet_cond', 'conditioning_scale', 'class_labels']
output_names=[]
input_names=['x', 'timesteps', 'controlnet_cond', 'conditioning_scale', 'context', 'class_labels']
export args: {'dynamo': False, 'report': True, 'dynamic_axes': {'x': [2, 3, 4]}}
�[0;93m2025-02-18 13:48:06.299308972 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /middle_block/attention/Unsqueeze_6�[m
�[0;93m2025-02-18 13:48:06.299401531 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /middle_block/attention/Unsqueeze_5�[m
�[0;93m2025-02-18 13:48:06.299450121 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /middle_block/attention/Unsqueeze_4�[m
�[0;93m2025-02-18 13:48:06.299485290 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /middle_block/attention/Unsqueeze_3�[m
�[0;93m2025-02-18 13:48:06.299519300 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /middle_block/attention/Unsqueeze_2�[m
�[0;93m2025-02-18 13:48:06.299561439 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /middle_block/attention/attn/out_rearrange/Unsqueeze_1�[m
�[0;93m2025-02-18 13:48:06.299596209 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /middle_block/attention/attn/out_rearrange/Unsqueeze�[m
�[0;93m2025-02-18 13:48:06.299634008 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /middle_block/attention/attn/input_rearrange_2/Unsqueeze_1�[m
�[0;93m2025-02-18 13:48:06.299672918 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /middle_block/attention/attn/input_rearrange_2/Unsqueeze�[m
�[0;93m2025-02-18 13:48:06.299710608 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /middle_block/attention/attn/input_rearrange_1/Unsqueeze_1�[m
�[0;93m2025-02-18 13:48:06.299748687 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /middle_block/attention/attn/input_rearrange_1/Unsqueeze�[m
�[0;93m2025-02-18 13:48:06.299786057 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /middle_block/attention/attn/input_rearrange/Unsqueeze_1�[m
�[0;93m2025-02-18 13:48:06.299825116 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /middle_block/attention/attn/input_rearrange/Unsqueeze�[m
�[0;93m2025-02-18 13:48:06.299864096 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /middle_block/attention/Unsqueeze_1�[m
�[0;93m2025-02-18 13:48:06.299902995 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /middle_block/attention/Unsqueeze�[m
�[0;93m2025-02-18 13:48:06.299941745 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.3/attentions.1/Unsqueeze_6�[m
�[0;93m2025-02-18 13:48:06.299979654 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.3/attentions.1/Unsqueeze_5�[m
�[0;93m2025-02-18 13:48:06.300016774 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.3/attentions.1/Unsqueeze_4�[m
�[0;93m2025-02-18 13:48:06.300053684 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.3/attentions.1/Unsqueeze_3�[m
�[0;93m2025-02-18 13:48:06.300111503 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.3/attentions.1/Unsqueeze_2�[m
�[0;93m2025-02-18 13:48:06.300151302 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.3/attentions.1/attn/out_rearrange/Unsqueeze_1�[m
�[0;93m2025-02-18 13:48:06.300189872 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.3/attentions.1/attn/out_rearrange/Unsqueeze�[m
�[0;93m2025-02-18 13:48:06.300229722 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.3/attentions.1/attn/input_rearrange_2/Unsqueeze_1�[m
�[0;93m2025-02-18 13:48:06.300274131 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.3/attentions.1/attn/input_rearrange_2/Unsqueeze�[m
�[0;93m2025-02-18 13:48:06.300314331 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.3/attentions.1/attn/input_rearrange_1/Unsqueeze_1�[m
�[0;93m2025-02-18 13:48:06.300351900 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.3/attentions.1/attn/input_rearrange_1/Unsqueeze�[m
�[0;93m2025-02-18 13:48:06.300391480 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.3/attentions.1/attn/input_rearrange/Unsqueeze_1�[m
�[0;93m2025-02-18 13:48:06.300431079 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.3/attentions.1/attn/input_rearrange/Unsqueeze�[m
�[0;93m2025-02-18 13:48:06.300470049 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.3/attentions.1/Unsqueeze_1�[m
�[0;93m2025-02-18 13:48:06.300522238 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.3/attentions.1/Unsqueeze�[m
�[0;93m2025-02-18 13:48:06.300565658 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.3/attentions.0/Unsqueeze_6�[m
�[0;93m2025-02-18 13:48:06.300603807 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.3/attentions.0/Unsqueeze_5�[m
�[0;93m2025-02-18 13:48:06.300640867 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.3/attentions.0/Unsqueeze_4�[m
�[0;93m2025-02-18 13:48:06.300678406 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.3/attentions.0/Unsqueeze_3�[m
�[0;93m2025-02-18 13:48:06.300713306 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.3/attentions.0/Unsqueeze_2�[m
�[0;93m2025-02-18 13:48:06.300753615 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.3/attentions.0/attn/out_rearrange/Unsqueeze_1�[m
�[0;93m2025-02-18 13:48:06.300791095 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.3/attentions.0/attn/out_rearrange/Unsqueeze�[m
�[0;93m2025-02-18 13:48:06.300836325 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.3/attentions.0/attn/input_rearrange_2/Unsqueeze_1�[m
�[0;93m2025-02-18 13:48:06.300876614 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.3/attentions.0/attn/input_rearrange_2/Unsqueeze�[m
�[0;93m2025-02-18 13:48:06.300916594 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.3/attentions.0/attn/input_rearrange_1/Unsqueeze_1�[m
�[0;93m2025-02-18 13:48:06.300954303 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.3/attentions.0/attn/input_rearrange_1/Unsqueeze�[m
�[0;93m2025-02-18 13:48:06.301003743 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.3/attentions.0/attn/input_rearrange/Unsqueeze_1�[m
�[0;93m2025-02-18 13:48:06.301041992 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.3/attentions.0/attn/input_rearrange/Unsqueeze�[m
�[0;93m2025-02-18 13:48:06.301080902 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.3/attentions.0/Unsqueeze_1�[m
�[0;93m2025-02-18 13:48:06.301118811 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.3/attentions.0/Unsqueeze�[m
�[0;93m2025-02-18 13:48:06.301156011 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.2/attentions.1/Unsqueeze_6�[m
�[0;93m2025-02-18 13:48:06.301193270 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.2/attentions.1/Unsqueeze_5�[m
�[0;93m2025-02-18 13:48:06.301230290 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.2/attentions.1/Unsqueeze_4�[m
�[0;93m2025-02-18 13:48:06.301267649 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.2/attentions.1/Unsqueeze_3�[m
�[0;93m2025-02-18 13:48:06.301307619 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.2/attentions.1/Unsqueeze_2�[m
�[0;93m2025-02-18 13:48:06.301343669 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.2/attentions.1/attn/out_rearrange/Unsqueeze_1�[m
�[0;93m2025-02-18 13:48:06.301377608 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.2/attentions.1/attn/out_rearrange/Unsqueeze�[m
�[0;93m2025-02-18 13:48:06.301412548 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.2/attentions.1/attn/input_rearrange_2/Unsqueeze_1�[m
�[0;93m2025-02-18 13:48:06.301445887 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.2/attentions.1/attn/input_rearrange_2/Unsqueeze�[m
�[0;93m2025-02-18 13:48:06.301480797 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.2/attentions.1/attn/input_rearrange_1/Unsqueeze_1�[m
�[0;93m2025-02-18 13:48:06.301514067 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.2/attentions.1/attn/input_rearrange_1/Unsqueeze�[m
�[0;93m2025-02-18 13:48:06.301548776 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.2/attentions.1/attn/input_rearrange/Unsqueeze_1�[m
�[0;93m2025-02-18 13:48:06.301582516 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.2/attentions.1/attn/input_rearrange/Unsqueeze�[m
�[0;93m2025-02-18 13:48:06.301616935 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.2/attentions.1/Unsqueeze_1�[m
�[0;93m2025-02-18 13:48:06.301652295 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.2/attentions.1/Unsqueeze�[m
�[0;93m2025-02-18 13:48:06.301685565 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.2/attentions.0/Unsqueeze_6�[m
�[0;93m2025-02-18 13:48:06.301718784 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.2/attentions.0/Unsqueeze_5�[m
�[0;93m2025-02-18 13:48:06.301752824 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.2/attentions.0/Unsqueeze_4�[m
�[0;93m2025-02-18 13:48:06.301786073 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.2/attentions.0/Unsqueeze_3�[m
�[0;93m2025-02-18 13:48:06.301827053 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.2/attentions.0/Unsqueeze_2�[m
�[0;93m2025-02-18 13:48:06.301864053 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.2/attentions.0/attn/out_rearrange/Unsqueeze_1�[m
�[0;93m2025-02-18 13:48:06.301897542 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.2/attentions.0/attn/out_rearrange/Unsqueeze�[m
�[0;93m2025-02-18 13:48:06.301932392 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.2/attentions.0/attn/input_rearrange_2/Unsqueeze_1�[m
�[0;93m2025-02-18 13:48:06.301965431 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.2/attentions.0/attn/input_rearrange_2/Unsqueeze�[m
�[0;93m2025-02-18 13:48:06.302001221 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.2/attentions.0/attn/input_rearrange_1/Unsqueeze_1�[m
�[0;93m2025-02-18 13:48:06.302035291 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.2/attentions.0/attn/input_rearrange_1/Unsqueeze�[m
�[0;93m2025-02-18 13:48:06.302070190 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.2/attentions.0/attn/input_rearrange/Unsqueeze_1�[m
�[0;93m2025-02-18 13:48:06.302102910 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.2/attentions.0/attn/input_rearrange/Unsqueeze�[m
�[0;93m2025-02-18 13:48:06.302136289 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.2/attentions.0/Unsqueeze_1�[m
�[0;93m2025-02-18 13:48:06.302169529 [W:onnxruntime:, unsqueeze_elimination.cc:20 Apply] UnsqueezeElimination cannot remove node /down_blocks.2/attentions.0/Unsqueeze�[m
torch_versioned_kwargs={'dynamo': False, 'report': True}
[I] Folding Constants | Pass 1
[I]     Total Nodes | Original:  1376, After Folding:   906 |   470 Nodes Folded
[I] Folding Constants | Pass 2
[I]     Total Nodes | Original:   906, After Folding:   515 |   391 Nodes Folded
[I] Folding Constants | Pass 3
[I]     Total Nodes | Original:   515, After Folding:   515 |     0 Nodes Folded
[I] Saving ONNX model to: /tmp/tmpza391yxe/model.onnx
2025-02-18 13:48:09,336 - INFO - Export to ONNX successful.
2025-02-18 13:48:09,337 - INFO - Building TensorRT engine for /tmp/tmpza391yxe/model.onnx: /tmp.plan
[!] Invalid Engine. Please ensure the engine was built correctly
[W] ModelImporter.cpp:420: Make sure input timesteps has Int64 binding.
[W] ModelImporter.cpp:420: Make sure input class_labels has Int64 binding.
[I] Configuring with profiles:[
        Profile 0:
            {x [min=[1, 4, 64, 64, 32], opt=[1, 4, 128, 128, 128], max=[1, 4, 128, 128, 128]],
             timesteps [min=[1], opt=[1], max=[1]],
             controlnet_cond [min=[1, 8, 256, 256, 256], opt=[1, 8, 256, 256, 256], max=[1, 8, 256, 256, 256]],
             conditioning_scale [min=[], opt=[], max=[]],
             class_labels [min=[1], opt=[1], max=[1]]}
    ]
[W] profileSharing0806 is on by default in TensorRT 10.0. This flag is deprecated and has no effect.
[I] Building engine with configuration:
    Flags                  | [FP16, TF32, OBEY_PRECISION_CONSTRAINTS]
    Engine Capability      | EngineCapability.STANDARD
    Memory Pools           | [WORKSPACE: 81153.75 MiB, TACTIC_DRAM: 81153.75 MiB, TACTIC_SHARED_MEMORY: 1024.00 MiB]
    Tactic Sources         | [EDGE_MASK_CONVOLUTIONS, JIT_CONVOLUTIONS]
    Profiling Verbosity    | ProfilingVerbosity.DETAILED
    Preview Features       | [PROFILE_SHARING_0806]
[E] IBuilder::buildSerializedNetwork: Error Code 4: API Usage Error (Input tensor x has static dimensions that don't match kMIN dimensions in profile index 0. Input dimensions are [1,4,64,64,64] but profile dimensions are [1,4,64,64,32].)
Traceback (most recent call last):
  File "/home/liubin/maisi/test/test_controlnet_trt.py", line 95, in <module>
    test()
  File "/home/liubin/maisi/test/test_controlnet_trt.py", line 92, in test
    output = trt_control_net(input_noise, timesteps=timesteps, controlnet_cond=controlnet_cond, class_labels=torch.randint(1,2,(1,),device=torch.device("cuda")))
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/monai/monai/networks/trt_compiler.py", line 599, in trt_forward
    return self._trt_compiler.forward(self, argv, kwargs)
  File "/opt/monai/monai/networks/trt_compiler.py", line 451, in forward
    raise e
  File "/opt/monai/monai/networks/trt_compiler.py", line 442, in forward
    self._build_and_save(model, build_args)
  File "/opt/monai/monai/networks/trt_compiler.py", line 588, in _build_and_save
    engine_bytes = self._onnx_to_trt(onnx_path)
  File "/opt/monai/monai/networks/trt_compiler.py", line 508, in _onnx_to_trt
    return engine_bytes_from_network(network, config=CreateConfig(profiles=profiles, **build_args))
  File "<string>", line 3, in engine_bytes_from_network
  File "/usr/local/lib/python3.10/dist-packages/polygraphy/backend/base/loader.py", line 40, in __call__
    return self.call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/polygraphy/util/util.py", line 733, in wrapped
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/polygraphy/backend/trt/loader.py", line 585, in call_impl
    G_LOGGER.critical(
  File "/usr/local/lib/python3.10/dist-packages/polygraphy/logger/logger.py", line 609, in critical
    raise ExceptionType(message) from None
polygraphy.exception.exception.PolygraphyException: Invalid Engine. Please ensure the engine was built correctly

Environment

MONAI 1.4.1rc1
Ubuntu 20.04

Additional context
I think the root cause of this problem is that the ONNX model is generated using the tracing way which records a fix input size. So if we try to pass profiles with dyanmic shape inputs it won't work.
Also tried to use the dynamo/script export to export the model and the model cannot be exported to onnx through these ways.

[binliunls]

@borisfom for discussion

[binliunls]

Also the autoencoder has some issues with 512x512x512 input when try to export as TRT.
Reproduce:

1. Start MONAI 1.4.1rc1 container
2. Download the MAISI model from model zoo new version PR
3. Replace the inference_trt.json file with content:

{
    "+imports": [
        "$from monai.networks import trt_compile"
    ],
    "autoencoder": "$trt_compile(@autoencoder_def, @trained_autoencoder_path, submodule=['decoder'])"
}

Error log:

2025-02-19 04:29:53,234 - INFO - Export to ONNX successful.                                                              
2025-02-19 04:29:53,234 - INFO - Building TensorRT engine for /tmp/tmpvyyqw45z/model.onnx: ./models/autoencoder_epoch273.
pt.decoder.plan                                                                                                          
[E] ITensor::getDimensions: Error Code 4: API Usage Error (/blocks.6/Resize_output: tensor volume exceeds 2147483648, dimensions are [1,128,320,320,320])                                                                                         
[E] ModelImporter.cpp:937: While parsing node number 3421 [Resize -> "/blocks.6/Resize_output_0"]:                       
[E] ModelImporter.cpp:938: --- Begin node ---                                                                            
    input: "/blocks.5/Cast_output_0"                                                                                     
    input: ""                                                                                                            
    input: "/blocks.6/Constant_output_0"                                                                                 
    output: "/blocks.6/Resize_output_0"                                                                                  
    name: "/blocks.6/Resize"                                                                                             
    op_type: "Resize"                                                                                                    
    attribute {                                                                                                          
      name: "coordinate_transformation_mode"                                                                             
      s: "half_pixel"                                                                                                    
      type: STRING                                                                                                       
    }                                                                                                                    
    attribute {                                                                                                          
      name: "cubic_coeff_a"                                                                                              
      f: -0.75                                                                                                           
      type: FLOAT                                                                                                        
    }                                                                                                                    
    attribute {                                                                                                          
      name: "mode"                                                                                                       
      s: "linear"                                                                                                        
      type: STRING                                                                                                       
    }                                                                                                                    
    attribute {                                                                                                          
      name: "nearest_mode"                                                                                               
      s: "floor"                                                                                                         
      type: STRING                                                                                                       
    }                                                                                                                    
[E] ModelImporter.cpp:939: --- End node ---    
[E] ModelImporter.cpp:942: ERROR: ModelImporter.cpp:197 In function parseNode:
    [6] Invalid Node - /blocks.6/Resize
    ITensor::getDimensions: Error Code 4: API Usage Error (/blocks.6/Resize_output: tensor volume exceeds 2147483648, dim
ensions are [1,128,320,320,320])
[E] In node 3421 with name: /blocks.6/Resize and operator: Resize (parseNode): INVALID_NODE: Invalid Node - /blocks.6/Res
ize
    ITensor::getDimensions: Error Code 4: API Usage Error (/blocks.6/Resize_output: tensor volume exceeds 2147483648, dim
ensions are [1,128,320,320,320])
[!] Could not parse ONNX correctly

It looks like that the 512x512x512 size images caused some intermediate results larger than the size of resize node.

[borisfom]

This unfortunately looks like a hard TensorRT limitation of 2G elements in tensor:
[E] ITensor::getDimensions: Error Code 4: API Usage Error (/blocks.6/Resize_output: tensor volume exceeds 2147483648, dimensions are [1,128,320,320,320])

[binliunls]

This unfortunately looks like a hard TensorRT limitation of 2G elements in tensor: [E] ITensor::getDimensions: Error Code 4: API Usage Error (/blocks.6/Resize_output: tensor volume exceeds 2147483648, dimensions are [1,128,320,320,320])

Any chance we can work together with the TensorRT team about this one in the future release?