How to export a PyTorch model to TorchScript

TorchScript artifacts still sit at the handoff between PyTorch and legacy C++ services or deployment stacks, even though new compiler work belongs on torch.export. When a consumer explicitly requires TorchScript, a saved module carries the parameters and forward graph without depending on the original Python class.

Tracing fits inference modules whose tensor operations do not change with input data. torch.jit.trace records only the path executed for representative inputs and fixes training-sensitive behavior to the mode used during tracing, so data-dependent branches require torch.jit.script instead.

Only load TorchScript files from trusted, untampered sources because deserialization can execute malicious pickle data. Keep producer and consumer PyTorch versions compatible because an archive saved by a newer release may not load in an older runtime.

Steps to export a PyTorch model to TorchScript:

Define the ScoreModel inference module in export_torchscript.py.

export_torchscript.py

from pathlib import Path
 
import torch
from torch import nn
 
 
class ScoreModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(4, 8),
            nn.ReLU(),
            nn.Linear(8, 2),
        )
 
    def forward(self, features):
        return torch.softmax(self.layers(features), dim=-1)

Append deterministic trace inputs that match the model's feature shape to export_torchscript.py.
export_torchscript.py
```
torch.manual_seed(7)
model = ScoreModel().eval()
trace_input = torch.randn(2, 4)
check_inputs = [(torch.randn(1, 4),), (torch.randn(5, 4),)]
artifact = Path("score_model_torchscript.pt")
```
The sample uses four input features across three batch sizes. The module remains in eval() because tracing preserves the training or evaluation behavior present during capture.
Append TorchScript tracing and persistence to export_torchscript.py.
export_torchscript.py
```
with torch.inference_mode():
    traced = torch.jit.trace(
        model,
        trace_input,
        check_inputs=check_inputs,
    )
 
torch.jit.save(traced, artifact)
```
The additional check_inputs make the tracer compare the captured graph with eager execution at two more batch sizes. They do not make data-dependent Python control flow safe to trace.

Append reload comparison and failure handling to export_torchscript.py.

export_torchscript.py

loaded = torch.jit.load(artifact, map_location="cpu")
test_input = torch.randn(3, 4)
 
with torch.inference_mode():
    eager_output = model(test_input)
    loaded_output = loaded(test_input)
 
outputs_match = torch.allclose(
    loaded_output,
    eager_output,
    rtol=1e-5,
    atol=1e-6,
)
 
print(f"torch version: {torch.__version__}")
print(f"artifact: {artifact} ({artifact.stat().st_size} bytes)")
print(f"loaded output shape: {tuple(loaded_output.shape)}")
print(f"outputs match eager: {outputs_match}")
 
if not outputs_match:
    raise RuntimeError("The loaded TorchScript output differs from eager mode")

Run the completed export script with the project environment's Python executable.
```
$ .venv/bin/python export_torchscript.py
torch version: 2.13.0+cpu
artifact: score_model_torchscript.pt (8985 bytes)
loaded output shape: (3, 2)
outputs match eager: True
```
The displayed .venv/bin/python path represents the environment that contains the trained model and PyTorch.

Create run_torchscript.py only when the .pt file came from a trusted build or artifact store.

run_torchscript.py

import torch
 
model = torch.jit.load("score_model_torchscript.pt", map_location="cpu")
 
with torch.inference_mode():
    output = model(torch.ones(1, 4))
 
print(f"output shape: {tuple(output.shape)}")
print(f"probability sum: {output.sum().item():.6f}")

TorchScript deserialization can execute malicious pickle data from a tampered file.

Run the independent loader to confirm the TorchScript artifact executes without the original model class.
```
$ .venv/bin/python run_torchscript.py
output shape: (1, 2)
probability sum: 1.000000
```