How to export a PyTorch model with torch.export

Model export turns an eager PyTorch module into an ahead-of-time graph that compiler, deployment, and conversion tooling can inspect without re-running the original Python program. torch.export captures a traceable tensor program as an ExportedProgram when the model can be represented from representative inputs.

torch.export.export records the tensor computation from an nn.Module and the constraints attached to its inputs. Example tensors are part of that contract, so dimensions stay specialized unless a Dim declaration marks them as dynamic.

Use this path after the model is in evaluation mode and the forward method can run without unsupported Python side effects. A saved .pt2 artifact can be loaded again with torch.export.load in Python, but PyTorch marks saved export files as under active development, so keep producer and consumer PyTorch versions aligned for deployment handoffs.

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

   export_torch_export.py

   from pathlib import Path
    
   import torch
   from torch import nn
   from torch.export import Dim
    
    
   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_inputs = (torch.randn(2, 4),)
   dynamic_shapes = {"features": {0: Dim("batch", min=1, max=8)}}
   eager_output = model(*example_inputs)
    
   exported = torch.export.export(
       model,
       example_inputs,
       dynamic_shapes=dynamic_shapes,
   )
    
   test_input = torch.randn(3, 4)
   exported_output = exported.module()(test_input)
    
   artifact = Path("score_model.pt2")
   torch.export.save(exported, artifact)
   loaded = torch.export.load(artifact)
   loaded_output = loaded.module()(torch.randn(1, 4))
    
   print(f"torch version: {torch.__version__}")
   print(f"eager output shape: {tuple(eager_output.shape)}")
   print(f"exported type: {type(exported).__name__}")
   print(f"graph nodes: {len(list(exported.graph_module.graph.nodes))}")
   print(f"range constraints: {exported.range_constraints}")
   print(f"exported output shape: {tuple(exported_output.shape)}")
   print(f"artifact: {artifact} ({artifact.stat().st_size} bytes)")
   print(f"loaded output shape: {tuple(loaded_output.shape)}")

   Replace ScoreModel and example_inputs with the real module and inputs. Keep model.eval() for inference exports so dropout and batch-normalization layers use evaluation behavior.

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

   $ .venv/bin/python export_torch_export.py
   torch version: 2.12.1+cpu
   eager output shape: (2, 2)
   exported type: ExportedProgram
   graph nodes: 10
   range constraints: {s97: VR[1, 8]}
   exported output shape: (3, 2)
   artifact: score_model.pt2 (13605 bytes)
   loaded output shape: (1, 2)

   Replace .venv/bin/python with the Python executable for the active project environment. The symbolic name in range constraints can vary; the important part is the accepted batch range.

   Load .pt2 files only from trusted sources. torch.export.load uses Python pickle during deserialization.

3. Test the saved .pt2 artifact as a readable archive.

   $ .venv/bin/python -m zipfile --test score_model.pt2
   Done testing

   The archive test confirms the file is readable, while the loaded output shape from the previous step confirms the saved ExportedProgram can execute after loading.