TorchScript

TorchScript converts a PyTorch model into a serialized, self-contained representation that can run independently of Python — enabling deployment in C++ environments, mobile apps, and other settings where a full Python runtime isn't available or desirable.

The Two Conversion Approaches

Code — Tracing

import torch

model.eval()
example_input = torch.randn(1, 3, 224, 224)

traced_model = torch.jit.trace(model, example_input)
traced_model.save("model_traced.pt")

# Reload -- no Python model class definition needed
loaded = torch.jit.load("model_traced.pt")
output = loaded(example_input)

Code — Scripting (for Models with Control Flow)

class ConditionalModel(torch.nn.Module):
    def forward(self, x):
        if x.sum() > 0:              # data-dependent control flow
            return self.branch_a(x)
        else:
            return self.branch_b(x)

model = ConditionalModel()
scripted_model = torch.jit.script(model)   # tracing would silently miss the branch not taken
scripted_model.save("model_scripted.pt")

Tracing's key limitation: it only records the specific path of operations executed for the one example input provided — any conditional branch not taken during that trace is permanently missing from the traced graph, silently producing incorrect behavior for inputs that would have taken the other branch. Scripting avoids this by analyzing the actual code structure, correctly capturing all branches.

Why This Matters for Deployment

A TorchScript model is a self-contained, serialized artifact that can be loaded and run from C++ (via LibTorch) or other environments without requiring a Python interpreter at all — valuable for latency-sensitive production serving, mobile/edge deployment, or any environment where shipping a full Python dependency stack is impractical.

Common Mistakes

• Using torch.jit.trace on a model with genuine data-dependent control flow — the trace will silently capture only the branch taken for the specific example input, producing incorrect results for other inputs without any error or warning.
• Forgetting to call model.eval() before tracing — tracing while in training mode can bake in Dropout/BatchNorm training-mode behavior incorrectly into the resulting traced graph.

Interview Relevance

Q: "What's the key difference between torch.jit.trace and torch.jit.script, and when would you need scripting specifically?" Tracing runs the model once on example input and records only the exact operations executed for that specific run — any data-dependent conditional branch not taken during tracing is silently missing afterward. Scripting instead directly compiles the model's actual Python code, correctly capturing all control flow paths (conditionals, loops). Scripting is necessary whenever a model's forward pass includes behavior that depends on the actual input values, not just its shape.

Practice Question

You trace a model using an example input where a certain if-branch is never triggered. What happens when the traced model later receives an input that should trigger that branch?

Related DL Notes

• Model Serialization
• ONNX
• Model Optimization (Deployment)