Saving models compiled with Torch-TensorRT#

Note

torch.compile(backend="torch_tensorrt") uses JIT compilation — engines are built on first call and cannot be saved directly. To serialize a TRT engine to disk, use the AOT path: torch_tensorrt.compile(model, ir="dynamo", ...) followed by torch_tensorrt.save(). See Saving torch.compile models below for details.

Saving models compiled with Torch-TensorRT can be done using torch_tensorrt.save API.

Dynamo IR#

The output type of ir=dynamo compilation of Torch-TensorRT is torch.fx.GraphModule object by default. We can save this object in either TorchScript (torch.jit.ScriptModule), ExportedProgram (torch.export.ExportedProgram) or PT2 formats by specifying the output_format flag. Here are the options output_format will accept

  • exported_program : This is the default. We perform transformations on the graphmodule first and use torch.export.save to save the module.

  • torchscript : We trace the graphmodule via torch.jit.trace and save it via torch.jit.save.

  • PT2 Format : This is a next generation runtime for PyTorch models, allowing them to run in Python and in C++

  • executorch : We lower the graphmodule to an ExecuTorch .pte program, delegating the TensorRT engines to the ExecuTorch backend. Linux-only; requires the executorch package.

a) ExportedProgram#

Here’s an example usage

import torch
import torch_tensorrt

model = MyModel().eval().cuda()
inputs = [torch.randn((1, 3, 224, 224)).cuda()]
# trt_ep is a torch.fx.GraphModule object
trt_gm = torch_tensorrt.compile(model, ir="dynamo", arg_inputs=inputs)
torch_tensorrt.save(trt_gm, "trt.ep", arg_inputs=inputs)

# Later, you can load it and run inference
model = torch.export.load("trt.ep").module()
model(*inputs)

Saving Models with Dynamic Shapes#

When saving models compiled with dynamic shapes, you have two methods to preserve the dynamic shape specifications:

Method 1: Using torch.export.Dim (explicit)

Provide explicit dynamic_shapes parameter following torch.export’s pattern:

import torch
import torch_tensorrt

model = MyModel().eval().cuda()
example_input = torch.randn((2, 3, 224, 224)).cuda()

# Define dynamic batch dimension
dyn_batch = torch.export.Dim("batch", min=1, max=32)
dynamic_shapes = {"x": {0: dyn_batch}}

# Export with dynamic shapes
exp_program = torch.export.export(
    model, (example_input,),
    dynamic_shapes=dynamic_shapes,
    strict=False
)

# Compile with dynamic input specifications
trt_gm = torch_tensorrt.dynamo.compile(
    exp_program,
    arg_inputs=[torch_tensorrt.Input(
        min_shape=(1, 3, 224, 224),
        opt_shape=(8, 3, 224, 224),
        max_shape=(32, 3, 224, 224),
    )]
)

# Save with dynamic_shapes to preserve dynamic behavior
torch_tensorrt.save(
    trt_gm,
    "trt_dynamic.ep",
    arg_inputs=[example_input],
    dynamic_shapes=dynamic_shapes  # Same as used during export
)

# Load and use with different batch sizes
loaded_model = torch_tensorrt.load("trt_dynamic.ep").module()
output_bs4 = loaded_model(torch.randn(4, 3, 224, 224).cuda())
output_bs16 = loaded_model(torch.randn(16, 3, 224, 224).cuda())

Method 2: Using torch_tensorrt.Input

Pass torch_tensorrt.Input objects with min/opt/max shapes directly, and the dynamic shapes will be inferred automatically:

import torch
import torch_tensorrt

model = MyModel().eval().cuda()

# Define Input with dynamic shapes
inputs = [
    torch_tensorrt.Input(
        min_shape=(1, 3, 224, 224),
        opt_shape=(8, 3, 224, 224),
        max_shape=(32, 3, 224, 224),
        dtype=torch.float32,
        name="x"  # Optional: provides better dimension naming
    )
]

# Compile with Torch-TensorRT
trt_gm = torch_tensorrt.compile(model, ir="dynamo", arg_inputs=inputs)

# Save with Input objects - dynamic_shapes inferred automatically!
torch_tensorrt.save(
    trt_gm,
    "trt_dynamic.ep",
    arg_inputs=inputs  # Dynamic shapes inferred from Input objects
)

# Load and use with different batch sizes
loaded_model = torch_tensorrt.load("trt_dynamic.ep").module()
output_bs4 = loaded_model(torch.randn(4, 3, 224, 224).cuda())
output_bs16 = loaded_model(torch.randn(16, 3, 224, 224).cuda())

b) Torchscript#

import torch
import torch_tensorrt

model = MyModel().eval().cuda()
inputs = [torch.randn((1, 3, 224, 224)).cuda()]
# trt_gm is a torch.fx.GraphModule object
trt_gm = torch_tensorrt.compile(model, ir="dynamo", arg_inputs=inputs)
torch_tensorrt.save(trt_gm, "trt.ts", output_format="torchscript", arg_inputs=inputs)

# Later, you can load it and run inference
model = torch.jit.load("trt.ts").cuda()
model(*inputs)

Torchscript IR#

In Torch-TensorRT 1.X versions, the primary way to compile and run inference with Torch-TensorRT is using Torchscript IR. For ir=ts, this behavior stays the same in 2.X versions as well.

import torch
import torch_tensorrt

model = MyModel().eval().cuda()
inputs = [torch.randn((1, 3, 224, 224)).cuda()]
trt_ts = torch_tensorrt.compile(model, ir="ts", arg_inputs=inputs) # Output is a ScriptModule object
torch.jit.save(trt_ts, "trt_model.ts")

# Later, you can load it and run inference
model = torch.jit.load("trt_model.ts").cuda()
model(*inputs)

Loading the models#

We can load torchscript or exported_program models using torch.jit.load and torch.export.load APIs from PyTorch directly. Alternatively, we provide a light wrapper torch_tensorrt.load(file_path) which can load either of the above model types.

Here’s an example usage

import torch
import torch_tensorrt

# file_path can be trt.ep or trt.ts file obtained via saving the model (refer to the above section)
inputs = [torch.randn((1, 3, 224, 224)).cuda()]
model = torch_tensorrt.load(<file_path>).module()
model(*inputs)

b) PT2 Format (AOTInductor)#

PT2 packages the model using AOTInductor, which compiles non-TRT subgraphs into native CUDA kernels. The resulting .pt2 file can be loaded in Python or C++ without a Torch-TensorRT runtime dependency.

import torch
import torch_tensorrt

model = MyModel().eval().cuda()
inputs = [torch.randn((1, 3, 224, 224)).cuda()]
trt_gm = torch_tensorrt.compile(model, ir="dynamo", arg_inputs=inputs)
torch_tensorrt.save(trt_gm, "trt.pt2", arg_inputs=inputs,
                    output_format="aot_inductor", retrace=True)

# Load without Torch-TensorRT at inference time
model = torch._inductor.aoti_load_package("trt.pt2")
model(*inputs)

For dynamic shapes, C++ deployment, and a full comparison with the ExportedProgram format, see AOTInductor Deployment.

c) ExecuTorch (.pte)#

The executorch output format lowers the compiled module to an ExecuTorch .pte program, delegating the TensorRT engines to the Torch-TensorRT ExecuTorch backend. It requires the executorch package (pip install "torch_tensorrt[executorch]") and is Linux-only.

import torch
import torch_tensorrt

model = MyModel().eval().cuda()
inputs = [torch.randn((1, 3, 224, 224)).cuda()]
trt_gm = torch_tensorrt.compile(model, ir="dynamo", arg_inputs=inputs)
torch_tensorrt.save(
    trt_gm, "trt.pte", output_format="executorch",
    retrace=False, arg_inputs=inputs,
)

Coalesced TensorRT + CUDA .pte

To run the ops TensorRT does not take on ExecuTorch’s CUDA (AOTInductor) backend instead of leaving them non-delegated, pass a CudaPartitioner via partitioners=. It is appended after the TensorRT partitioner, so TensorRT claims what it can and the CudaPartitioner picks up the rest as a catch-all:

from executorch.backends.cuda.cuda_backend import CudaBackend
from executorch.backends.cuda.cuda_partitioner import CudaPartitioner

torch_tensorrt.save(
    trt_gm, "trt.pte", output_format="executorch",
    retrace=False, arg_inputs=inputs,
    partitioners=[
        CudaPartitioner(
            [CudaBackend.generate_method_name_compile_spec("forward")]
        )
    ],
)

This needs ExecuTorch’s CUDA backend and a CUDA toolkit (nvcc/ptxas) at export time, and produces a .pte that requires a CUDA runtime at load. Any external CUDA weights are written as .ptd data file(s) next to the .pte; the runtime must be pointed at those data files to load them.

Warning

The CUDA backend names its external weight blob per-device (e.g. aoti_cuda_blob.ptd), not per-model, so saving two different coalesced .pte into the same directory overwrites the blob and the first .pte will fail to load. Save each coalesced model into its own directory.

Saving torch.compile models#

torch.compile(backend="torch_tensorrt") is a JIT path — TRT engines are built lazily on the first call and live in memory only. There is no direct way to call torch_tensorrt.save() on a torch.compile-compiled model.

Workaround: switch to the AOT path

Replace torch.compile with torch_tensorrt.compile(ir="dynamo") to get a serializable torch.fx.GraphModule:

import torch
import torch_tensorrt

model = MyModel().eval().cuda()
inputs = [torch.randn((1, 3, 224, 224)).cuda()]

# JIT path — NOT serializable
# jit_model = torch.compile(model, backend="torch_tensorrt", ...)

# AOT path — produces a serializable GraphModule
trt_gm = torch_tensorrt.compile(
    model,
    ir="dynamo",
    arg_inputs=inputs,
)

# Save to disk
torch_tensorrt.save(trt_gm, "model.ep", arg_inputs=inputs)

# Reload and run — no recompilation needed
loaded = torch_tensorrt.load("model.ep").module()
output = loaded(*inputs)

The ir="dynamo" path supports all the same compilation options as torch.compile(backend="torch_tensorrt"). The key differences are:

Feature

torch.compile (JIT)

ir="dynamo" (AOT)

Compilation timing

On first call

Explicit compile step

Auto-recompile on shape change

Yes

No (fixed shapes unless dynamic Input used)

Serializable to disk

No

Yes (torch_tensorrt.save)

C++ deployment

No

Yes (via ExportedProgram or PT2 format)
On this page
Edit on GitHub
Show Source