Checkpoints, Export and Inference
TorchScript: tracing vs scripting
Tracing records one run of your model and keeps whatever happened to happen; scripting reads your source code and keeps the logic — which is why tracing silently produces a wrong model whenever your forward pass has an if.
- 10 min read
- 3 reading levels
- Published
Read these first
On this page 6
One lesson, three depths. Pick the one that fits you today — you can switch any time.
Beginner — No maths. Plain English.
Two ways to turn a model into a standalone file: film one run of it, or translate its instructions.
Imagine explaining how to get to your house. Tracing is filming yourself walking there once, on a Tuesday. Anyone can follow the video. But if the film shows you turning left because that Tuesday the right-hand road was closed, the video now says "always turn left".
Scripting is writing the directions down: "if the road is closed, turn left, otherwise turn right." Harder to write, correct in every situation.
Why anybody does this
Your model needs your Python class to run. That is fine on your laptop and awkward everywhere else — a phone, a C++ server, a browser.
Turning the model into a self-contained file removes that dependency. The file carries both the numbers and the steps. Load it anywhere and it runs, with no model class and sometimes with no Python at all.
The trap
Tracing is one line of code and works on nearly any model, so people reach for it. It also fails silently.
Say your model makes a decision based on its input: if this value is large, do something different. The film records only the branch that ran on the day of filming. Feed it different data later and it takes the recorded turn regardless. No warning at the point of use. Wrong answers.
How it works
your model: if the input adds up to a positive number, double it
otherwise flip the sign
traced with a POSITIVE example:
recorded steps: double it <- the "if" is gone
feed a negative input -> still doubles it. WRONG.
scripted from the source code:
recorded steps: if positive: double, else: flip
feed a negative input -> flips. Correct.A real example you have seen
Route directions on a map app versus a screen recording of someone using it. The recording only helps if you start where they started, on the day they did. The directions adapt.
Remember this
- Tracing records one run; anything that could have gone another way is lost.
- Scripting reads your code and keeps the branches, at the cost of stricter rules.
- If your
forwardhas anifor awhilethat depends on the data, do not trace it.
What to learn next
- torch.export and the ExportedProgram — the modern replacement, and how it handles the same branch.
- ONNX — the other portable format, and the one most runtimes consume.
- torch.compile — a compiler for speed rather than for portability.
Developer — Code and libraries.
Setup
pip install torchRuns on CPU. First, a warning worth having up front: TorchScript is in maintenance mode. It is not going away, an enormous amount of deployed code uses it, and PyTorch is explicit that no new features are coming. For new work the successor is torch.export. Learn TorchScript because you will meet it, not because you should design around it.
Watching a trace go wrong
import torch
import torch.nn as nn
class Clamped(nn.Module):
def forward(self, x):
if x.sum() > 0: # a real branch: the answer depends on the data
return x * 2
return x * -1
model = Clamped()
positive = torch.ones(3)
negative = -torch.ones(3)
traced = torch.jit.trace(model, positive) # recorded with a POSITIVE example
scripted = torch.jit.script(model) # compiled from the source code
print("eager on negative input:", model(negative).tolist())
print("traced on negative input:", traced(negative).tolist())
print("scripted on negative input:", scripted(negative).tolist())
print()
print(scripted.code)trace_vs_script.py:6: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
if x.sum() > 0: # a real branch: the answer depends on the data
eager on negative input: [1.0, 1.0, 1.0]
traced on negative input: [-2.0, -2.0, -2.0]
scripted on negative input: [1.0, 1.0, 1.0]
def forward(self,
x: Tensor) -> Tensor:
if bool(torch.gt(torch.sum(x), 0)):
_0 = torch.mul(x, 2)
else:
_0 = torch.mul(x, -1)
return _0The correct answer for a negative input is 1.0. The traced model returns -2.0. It took the x * 2 branch it was filmed taking, on data where the model should have flipped the sign.
PyTorch did warn, and this is the important habit: TracerWarning is an error in disguise. Never let one through a code review.
The printed scripted.code is TorchScript's own view of your function, translated back into readable Python. The if survived. Printing .code after scripting is the fastest way to confirm the compiler understood what you meant.
The file runs without your class
import torch
import torch.nn as nn
class Clamped(nn.Module):
def forward(self, x):
if x.sum() > 0:
return x * 2
return x * -1
torch.jit.script(Clamped()).save("clamped.pt")
print("saved")import torch # note: the Clamped class is NOT defined here
model = torch.jit.load("clamped.pt")
model.eval()
print(model(torch.ones(3)).tolist())
print(model(-torch.ones(3)).tolist())saved [2.0, 2.0, 2.0] [1.0, 1.0, 1.0]
The second file has no class Clamped anywhere in it, and both branches still work. That is the entire selling point: the archive holds the weights and the logic, so the consumer needs neither your source file nor, from C++, a Python interpreter at all.
Choosing between them
torch.jit.trace | torch.jit.script | |
|---|---|---|
| how it works | runs the model once, records tensor operations | compiles the Python source |
data-dependent if / while | silently baked in | preserved |
| Python that is not tensor code | dropped, without notice | must be TorchScript-compatible or it fails |
| variable input shapes | usually fine | fine |
| effort on a complex model | low | can be substantial |
| failure mode | wrong answers later | a compile error now |
The practical instinct: prefer scripting; a loud compile error today beats a silent wrong answer in production. Trace only when the module is pure tensor arithmetic with no branching, or when scripting refuses something you cannot restructure.
They also mix. @torch.jit.script a small branching helper, then trace the module that calls it — the traced graph calls into the scripted function and the logic survives.
What scripting refuses
TorchScript is a typed subset of Python. Common things it rejects:
- most third-party libraries inside
forward— NumPy, Pandas, PIL - a variable that holds an
intand later aTensor **kwargs, most decorators, and many comprehension formsOptional[Tensor]used without aNonecheck first- inheritance patterns beyond a plain
nn.Module
Mark helpers it should leave alone with @torch.jit.ignore (calling one from a scripted path fails at runtime) or @torch.jit.unused (replaced by a raise). Both are ways of saying "this branch never runs in production".
Common mistakes
Ignoring TracerWarning. It is the only notice you get.
Tracing a model whose forward reads self.training. The trace records the mode it was in. Call .eval() before tracing an inference artefact, or dropout stays active in your served model.
Tracing a model with a data-dependent loop count. The loop is unrolled to the length it happened to run. Same failure as the if, harder to spot.
Assuming TorchScript makes it faster. It removes Python overhead, which matters for small models and matters little for large ones. Measure — see timing GPU code correctly.
Scripting a model containing a torch.compiled submodule. They are different compilers and do not nest. Pick one.
Shipping a traced model without testing the other branch. The test that would have caught the bug above is two lines long.
Try it yourself
Add for _ in range(int(x.sum())): to Clamped.forward, then trace it with torch.ones(3) and call it with torch.ones(5). Count how many times the loop body appears in traced.code.
What to learn next
- torch.export and the ExportedProgram — the modern replacement, and how it handles the same branch.
- ONNX — the other portable format, and the one most runtimes consume.
- torch.compile — a compiler for speed rather than for portability.
Researcher — Mathematics and papers.
Two front ends, one IR
Both paths produce a ScriptModule wrapping a torch::jit::Graph in static single assignment form over a typed value system (Tensor, int, float, bool, List[T], Dict[K,V], Optional[T], Tuple[...], plus module and class types). They differ only in how the graph is built.
Tracing executes the module under a dispatcher-level recorder. Every ATen operator dispatched during that execution is appended to the graph; every Python-level operation — control flow, prints, list mutation, .item() comparisons — is invisible to the recorder and does not appear at all. Converting a tensor to a Python bool triggers the TracerWarning because the resulting branch decision becomes a compile-time constant. Shapes are recorded but not, by default, constrained, so a traced graph generalises across shapes and not across control flow.
Scripting parses the function's source with Python's ast module and lowers it to the same IR through a type-inference pass. Types are inferred from annotations and usage, and any construct outside the supported subset produces a compile error at torch.jit.script time. This is why scripting's failures are loud and tracing's are silent: one analyses the program, the other observes one execution of it.
The saved archive is a ZIP holding serialised bytecode, the constant pool, the parameter tensors, and the pickled module structure — self-contained, loadable by torch::jit::load in C++ with no Python runtime. That property, not speed, is TorchScript's reason for existing.
Optimisation, and why it is modest
The JIT applies profile-guided optimisation across a few executions: dead code elimination, constant propagation, common subexpression elimination, peephole rewrites, and pointwise fusion through NNC or nvFuser depending on version and device. Fusion of elementwise chains is the largest real win, particularly on GPU where it removes intermediate memory traffic. It does not restructure algorithms, does not change memory layout, and does not touch the matrix multiplications where most large-model time is spent — hence the observed pattern that TorchScript gives worthwhile speedups on small, launch-bound models and near-nothing on transformer-scale ones. The first few calls are slower while profiling runs, so benchmarks must discard warmup.
Status, and the successor
PyTorch documents TorchScript as no longer under active development. The successor stack is torch.export, which uses TorchDynamo's bytecode analysis to capture a graph, producing an ExportedProgram over ATen operators with explicit dynamic-shape constraints. The two differ in a way that matters for exactly the example above: torch.export refuses data-dependent control flow with a UserError naming the line, rather than specialising on it. Refusing is a design choice motivated by the tracing failure mode this lesson demonstrates — see the export lesson for the error and the cond operator that expresses the branch explicitly.
Existing TorchScript artefacts continue to load, and the C++ deployment path remains supported, so migration is a decision about new work rather than an urgent one.
References
- PyTorch documentation, TorchScript — the language reference, the supported Python subset, and the maintenance-mode notice.
- DeVito et al. (2022), TorchScript: Optimized Execution of PyTorch Programs — the IR, the two front ends, and the profiling executor.
- Ansel et al. (2024), PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation, ASPLOS — why bytecode analysis replaced both tracing and source parsing.
What to learn next
- torch.export and the ExportedProgram — the modern replacement, and how it handles the same branch.
- ONNX — the other portable format, and the one most runtimes consume.
- torch.compile — a compiler for speed rather than for portability.