Table of Contents
What TorchScript Is and Why It Exists
TorchScript is a way to turn a PyTorch model into a representation that can run without Python. This is useful when you want to ship a model into an environment where Python is inconvenient or unavailable, or when you want a more self contained artifact for production style execution. TorchScript models can be executed by the PyTorch runtime in C++ and other contexts that embed the runtime.
TorchScript is not required for every deployment. Many real systems use Python inference servers. TorchScript is mainly valuable when you need to reduce Python dependencies, you want to embed a model into a larger application, or you want a more constrained, serializable program representation.
Two Paths Into TorchScript: Scripting and Tracing
There are two main ways to produce TorchScript. Scripting analyzes your model code and compiles a TorchScript program that tries to preserve control flow such as if statements and loops. Tracing runs an example input through the model and records the operations that happened, producing a graph of those operations.
This chapter focuses on tracing basics, but it is important to know the core limitation of tracing. Tracing captures what happened for the given example inputs, not what could happen for different inputs.
Tracing records a single execution. If your model has input dependent control flow, tracing can produce an incorrect model because the recorded graph may not match other inputs.
A Minimal Tracing Workflow in PyTorch
The typical tracing workflow has four steps. You put the model in evaluation mode, you create example inputs that match real inference shapes and dtypes, you trace, and you save the traced artifact.
You usually want to use torch.no_grad() during tracing to avoid unnecessary autograd tracking, and you should call model.eval() so layers like dropout and batch normalization behave in inference mode.
A minimal example looks like this.
import torch
import torch.nn as nn
class SmallNet(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(10, 32),
nn.ReLU(),
nn.Linear(32, 3),
)
def forward(self, x):
return self.net(x)
model = SmallNet()
model.eval()
example = torch.randn(4, 10)
with torch.no_grad():
traced = torch.jit.trace(model, example)
traced.save("smallnet_traced.pt")The saved file contains the traced program plus the parameter values. You can load and run it like this.
loaded = torch.jit.load("smallnet_traced.pt")
loaded.eval()
x = torch.randn(2, 10)
y = loaded(x)Validating a Traced Model
A traced model can load and run but still be wrong if tracing missed behavior. A simple check is to compare the traced output against the original model output on a few random test inputs that reflect your real input distribution.
model.eval()
loaded.eval()
for _ in range(5):
x = torch.randn(8, 10)
with torch.no_grad():
y_ref = model(x)
y_ts = loaded(x)
max_diff = (y_ref - y_ts).abs().max().item()
print(max_diff)You should expect the maximum difference to be very small for pure feedforward models without randomness. If your model has numerically sensitive operations, small floating point differences can occur, but large discrepancies indicate a tracing issue or a mode mismatch.
Always compare outputs of the original model and the traced model on multiple representative inputs. Successful saving and loading does not guarantee correctness.
Common Tracing Pitfalls
A frequent pitfall is forgetting eval() before tracing. If dropout is active, the traced graph will bake in a particular dropout mask behavior for the example run, and it will not reflect correct inference. If batch normalization is in training mode, the trace can capture behavior that depends on batch statistics rather than running estimates.
Another pitfall is input shape assumptions. Tracing specializes to the operations observed, and certain shape dependent paths can behave unexpectedly if you later pass a different shape. Some models handle variable shapes fine, but you should test the shapes you plan to deploy.
A third pitfall is Python side logic. If your forward uses Python constructs that affect computation such as if x.sum() > 0: ... then tracing will only record the branch taken for the example input. The traced model will not contain the other branch.
Trace only models whose forward computation is effectively the same graph for all inputs you will use in deployment. If the graph changes with the input, prefer scripting.
Tracing with Multiple Inputs and Keyword Arguments
If your model forward takes multiple tensors, you trace with a tuple of example inputs.
class PairNet(nn.Module):
def forward(self, a, b):
return a @ b
m = PairNet().eval()
a = torch.randn(2, 3)
b = torch.randn(3, 4)
tr = torch.jit.trace(m, (a, b))For keyword arguments, the most robust approach for tracing is to arrange your forward to accept tensors positionally for the traced entry point. If you must use keyword arguments, you can wrap the model in a small module whose forward uses positional inputs and forwards them appropriately.
Inspecting the Traced Graph at a High Level
TorchScript lets you print the traced code and graph, which can help you confirm that the model captured the operations you expect.
print(traced.code)
print(traced.graph)You are not expected to understand every detail as a beginner, but this is a helpful debugging tool when a trace does not behave as expected.
Device and dtype Considerations During Tracing
You should trace using the same device type you plan to use for inference, or at least validate on that device type. A common pattern is to trace on CPU for portability, then run on CPU in production. If you want a CUDA optimized path, you can move the model and example input to CUDA before tracing.
device = torch.device("cuda")
model = model.to(device).eval()
example = example.to(device)
with torch.no_grad():
traced = torch.jit.trace(model, example)
Be consistent with dtype as well. If you plan to run with float32, trace with float32. If you plan to run with mixed precision or float16, that enters topics beyond tracing basics, but the key rule is that your example inputs should match deployment reality.
When to Prefer Scripting Over Tracing
If your model includes data dependent branching, loops over sequence length that depends on the input, or any logic where different inputs should produce different computation graphs, tracing is risky. In those cases, scripting is the safer approach because it attempts to compile the control flow itself rather than recording one run.
This chapter does not cover scripting details, but the decision rule is simple.
Use tracing for straight line tensor programs. Use scripting when correct behavior depends on control flow.
Practical Checklist Before You Save a TorchScript Artifact
You should ensure the model is in eval mode, use representative example inputs, validate outputs on several inputs, and save the traced module. If your model takes structured inputs, consider wrapping it so the traced forward has a stable, tensor only signature. Once saved, load it in a fresh process and run a quick inference test to confirm the artifact is self contained.