Export and save problem(torch-tensorrt)

Bug Description

When I used torch-tensorrt to complite, some errors came out suddenly, which made me confuse.

To Reproduce

Code for method 1:

import torch
import tensorrt
import torch_tensorrt
from torchvision.models import resnet50


if __name__ == '__main__':
    model = resnet50().eval().to('cuda')
    inputs = [torch.randn((1, 3, 224, 224)).cuda()]
    trt_gm = torch_tensorrt.compile(model, ir="dynamo", inputs=inputs)
    trt_traced_model = torch.jit.trace(trt_gm, inputs)
    torch.jit.save(trt_traced_model, "trt_model.ts")

Error:

INFO:torch_tensorrt.dynamo._compiler:Compilation Settings: CompilationSettings(precision=torch.float32, debug=False, workspace_size=0, min_block_size=5, torch_executed_ops=set(), 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, output_format='exported_program')

INFO:torch_tensorrt.dynamo.conversion._TRTInterpreter:TRT INetwork construction elapsed time: 0:00:00.567413
INFO:torch_tensorrt.dynamo.conversion._TRTInterpreter:Build TRT engine elapsed time: 0:00:10.713493
INFO:torch_tensorrt.dynamo.conversion._TRTInterpreter:TRT Engine uses: 7426048 bytes of Memory
/home/dan/anaconda3/envs/dust3r/lib/python3.11/site-packages/torch/_export/exported_program.py:333: UserWarning: Unable to execute the generated python source code from the graph. The graph module will no longer be directly callable, but you can still run the ExportedProgram, and if needed, you can run the graph module eagerly using torch.fx.Interpreter.
  warnings.warn(
Traceback (most recent call last):
  File "/media/dan/DATA/Processing/OFFICE/CodingFactory/PosWorkShop/dust3r/torchtrt_demo.py", line 15, in <module>
    trt_traced_model = torch.jit.trace(trt_gm, inputs)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dan/anaconda3/envs/dust3r/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 489, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/dan/anaconda3/envs/dust3r/lib/python3.11/site-packages/torch/_dynamo/external_utils.py", line 17, in inner
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/dan/anaconda3/envs/dust3r/lib/python3.11/site-packages/torch/jit/_trace.py", line 861, in trace
    name = _qualified_name(func)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/dan/anaconda3/envs/dust3r/lib/python3.11/site-packages/torch/_jit_internal.py", line 1163, in _qualified_name
    raise RuntimeError("Could not get name of python class object")
**RuntimeError: Could not get name of python class object**

However, when I tried : # Code for method 2:

import torch
import tensorrt
import torch_tensorrt
from torchvision.models import resnet50


if __name__ == '__main__':
    model = resnet50().eval().to('cuda:0')
    inputs = [torch.randn((1, 3, 224, 224)).cuda()]
    trt_gm = torch_tensorrt.compile(model, ir="dynamo", inputs=inputs)
    trt_exp_program = torch_tensorrt.dynamo.export(trt_gm, inputs, 'ep')
    torch.export.save(trt_exp_program, "trt_model.ep")

Error:

INFO:torch_tensorrt.dynamo._compiler:Compilation Settings: CompilationSettings(precision=torch.float32, debug=False, workspace_size=0, min_block_size=5, torch_executed_ops=set(), 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, output_format='exported_program')

INFO:torch_tensorrt.dynamo.conversion._TRTInterpreter:TRT INetwork construction elapsed time: 0:00:00.521946
INFO:torch_tensorrt.dynamo.conversion._TRTInterpreter:Build TRT engine elapsed time: 0:00:10.723753
INFO:torch_tensorrt.dynamo.conversion._TRTInterpreter:TRT Engine uses: 7325696 bytes of Memory

/home/dan/anaconda3/envs/dust3r/lib/python3.11/site-packages/torch/_export/exported_program.py:333: UserWarning: Unable to execute the generated python source code from the graph. The graph module will no longer be directly callable, but you can still run the ExportedProgram, and if needed, you can run the graph module eagerly using torch.fx.Interpreter.
  warnings.warn(
Traceback (most recent call last):
  File "/media/dan/DATA/Processing/OFFICE/CodingFactory/PosWorkShop/dust3r/torchtrt_demo.py", line 21, in <module>
    trt_exp_program = torch_tensorrt.dynamo.export(trt_gm, inputs, 'ep')
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dan/anaconda3/envs/dust3r/lib/python3.11/site-packages/torch_tensorrt/dynamo/_exporter.py", line 33, in export
    patched_module = transform(gm, inputs)
                     ^^^^^^^^^^^^^^^^^^^^^
  File "/home/dan/anaconda3/envs/dust3r/lib/python3.11/site-packages/torch_tensorrt/dynamo/_exporter.py", line 59, in transform
    _, outputs_map = partitioning.run_shape_analysis(gm, inputs)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dan/anaconda3/envs/dust3r/lib/python3.11/site-packages/torch_tensorrt/dynamo/partitioning/common.py", line 30, in run_shape_analysis
    for name, _ in parent_module.named_children():
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
**AttributeError: 'ExportedProgram' object has no attribute 'named_children'**

However, when I set torch_tensorrt.dynamo.export(trt_gm, inputs, 'ep') torch_tensorrt.dynamo.export(trt_gm, inputs, 'fx'), everything was fine…

Environment

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

  • Torch-TensorRT Version (e.g. 1.0.0): 2.2.0
  • PyTorch Version (e.g. 1.0): 2.2.0 stable
  • CPU Architecture: x64
  • OS (e.g., Linux): ubuntu22.04 LTS
  • How you installed PyTorch (conda, pip, libtorch, source): pip install torch torchvision torch-tensorrt
  • Build command you used (if compiling from source): pip install
  • Are you using local sources or building from archives:
  • Python version: 3.11
  • CUDA version: CUDA 12.1 and cudnn 8.9.1
  • GPU models and configuration: RTX 4060 laptop
  • Any other relevant information: None

Additional context

Code1 and Code2 are failed, and TensorRT/examples/torchtrt_runtime_example/network.py at main · pytorch/TensorRT · GitHub met the same error.
Could you tell me how to solve these errors?
Thank you in advance!