torch.export — The Modern Export System

The replacement for TorchScript

torch.export is PyTorch's modern model export system. It captures your model into a clean, standardized graph representation that can be serialized, deployed, quantized, or lowered to other backends (ONNX, ExecuTorch, CoreML). It replaces the older TorchScript (torch.jit.trace / torch.jit.script) for new projects.

Why a new export system

TorchScript was hard. jit.trace couldn't see Python control flow. jit.script required a restricted Python subset. The error messages were notoriously cryptic. torch.export uses the same dynamo-based graph capture as torch.compile, which means much better coverage and far cleaner errors.

The contract

You give torch.export.export a model and example inputs. It runs the model with those inputs, captures the resulting graph, and returns an ExportedProgram. The exported graph is fully serializable — save it to disk and load it without needing the original Python class definition.

What torch.export buys you

Backend portability — the same exported program lowers to ONNX, ExecuTorch (mobile), CoreML, TensorRT.
Quantization — modern quant techniques in torchao operate on exported programs.
Optimization — graph passes (constant folding, dead code elimination) operate on the exported representation.
Versioning — the export format is meant to be stable across PyTorch versions.

TorchScript — when (rarely) you'd still reach for it

For deploying to environments where torch.export's runtime isn't yet available (some legacy embedded paths). For new projects, default to torch.export.

Code

Export and load a model·python

import torch
import torch.nn as nn

class TinyMLP(nn.Module):
    def __init__(self): super().__init__(); self.fc = nn.Linear(10, 4)
    def forward(self, x): return self.fc(x)

model = TinyMLP().eval()
example = torch.randn(1, 10)

# Export
exported = torch.export.export(model, (example,))

# Save — no Python class definition needed on load
torch.export.save(exported, "/tmp/tiny.pt2")

# Load (in another process / machine)
loaded = torch.export.load("/tmp/tiny.pt2")
y = loaded.module()(example)        # call .module() to get a callable
print(y.shape)                       # torch.Size([1, 4])

Export with dynamic batch size·python

import torch
from torch.export import Dim

class TinyMLP(torch.nn.Module):
    def __init__(self): super().__init__(); self.fc = torch.nn.Linear(10, 4)
    def forward(self, x): return self.fc(x)

model = TinyMLP().eval()
example = torch.randn(1, 10)

# Tell torch.export that dim 0 (batch) is dynamic
batch = Dim("batch", min=1, max=128)
exported = torch.export.export(
    model, (example,), dynamic_shapes=({0: batch},),
)

# Now the same exported program can run on different batch sizes
loaded = torch.export.load
torch.export.save(exported, "/tmp/dyn.pt2")
loaded = torch.export.load("/tmp/dyn.pt2")
print(loaded.module()(torch.randn(32, 10)).shape)   # torch.Size([32, 4])

ONNX export — modern dynamo path·python

import torch

class TinyMLP(torch.nn.Module):
    def __init__(self): super().__init__(); self.fc = torch.nn.Linear(10, 4)
    def forward(self, x): return self.fc(x)

model = TinyMLP().eval()
example = torch.randn(1, 10)

# In PyTorch 2.x, torch.onnx.export defaults to dynamo=True (the modern path)
torch.onnx.export(
    model,
    (example,),
    "/tmp/tiny.onnx",
    input_names=["x"],
    output_names=["y"],
    dynamic_axes={"x": {0: "batch"}, "y": {0: "batch"}},
)

TorchScript — legacy reference·python

import torch

class TinyMLP(torch.nn.Module):
    def __init__(self): super().__init__(); self.fc = torch.nn.Linear(10, 4)
    def forward(self, x): return self.fc(x)

model = TinyMLP().eval()
example = torch.randn(1, 10)

# Method 1: tracing — records ops from this example
traced = torch.jit.trace(model, example)
traced.save("/tmp/tiny_traced.pt")

# Method 2: scripting — compiles Python (restricted subset)
scripted = torch.jit.script(model)
scripted.save("/tmp/tiny_scripted.pt")

# Both can be loaded without Python — but for new projects, prefer torch.export

Exercise

Export the same TinyMLP three ways: torch.export, torch.jit.trace, torch.jit.script. Save all three. Time loading + one inference call for each. Diff the file sizes. The torch.export and TorchScript variants will all work; pick which one you'd ship by reading the error messages if any of them fails on a tricky input.