How to export a PyTorch model to TorchScript

Legacy serving stacks, mobile paths, and C++ loaders may still require a TorchScript artifact even when new PyTorch export work belongs on torch.export or ONNX. Exporting through torch.jit turns an nn.Module into a saved module that can be loaded without the original Python class definition.

The simplest inference path uses torch.jit.trace with representative tensor inputs. Tracing records the tensor operations from one forward pass, so it fits models whose branch choices do not change based on input data; models with TorchScript-supported data-dependent branches need torch.jit.script instead.

Keep the model in eval() before export and run a reload smoke test before handing the .pt file to another runtime. The saved file should load with torch.jit.load, produce the expected output shape, and match the eager module on a fresh input.

1. Create the TorchScript export smoke script beside the model code.

   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)
    
    
   torch.manual_seed(7)
    
   model = ScoreModel().eval()
   example_input = torch.randn(2, 4)
    
   with torch.inference_mode():
       eager_output = model(example_input)
    
   scripted = torch.jit.trace(model, example_input)
   artifact = Path("score_model_torchscript.pt")
   scripted.save(artifact)
    
   loaded = torch.jit.load(artifact, map_location="cpu")
   test_input = torch.randn(3, 4)
    
   with torch.inference_mode():
       loaded_output = loaded(test_input)
       eager_check = model(test_input)
    
   print(f"torch version: {torch.__version__}")
   print("export method: torch.jit.trace")
   print(f"scripted type: {type(scripted).__name__}")
   print(f"eager output shape: {tuple(eager_output.shape)}")
   print(f"artifact: {artifact} ({artifact.stat().st_size} bytes)")
   print(f"loaded output shape: {tuple(loaded_output.shape)}")
   print(f"outputs match eager: {torch.allclose(loaded_output, eager_check)}")

   Replace ScoreModel and example_input with the real model and representative inputs. Keep torch.jit.trace only when one traced tensor path represents the inference behavior; switch to torch.jit.script for TorchScript-supported control flow that must change with the input.

2. Run the export script with the Python environment that has PyTorch installed.

   $ .venv/bin/python export_torchscript.py
   torch version: 2.12.1+cpu
   export method: torch.jit.trace
   scripted type: TopLevelTracedModule
   eager output shape: (2, 2)
   artifact: score_model_torchscript.pt (9049 bytes)
   loaded output shape: (3, 2)
   outputs match eager: True

   Replace .venv/bin/python with the Python executable for the active project environment. map_location=“cpu” makes the reload smoke test independent of a GPU on the validation host.

   Load TorchScript files only from trusted sources. torch.jit.load can execute malicious pickle data during deserialization.

3. Test the saved TorchScript archive for readability.

   $ .venv/bin/python -m zipfile --test score_model_torchscript.pt
   Done testing

   The archive test confirms that the saved file is readable, and the reload smoke output confirms executable behavior through loaded output shape and outputs match eager.