How to export a PyTorch model with torch.export

Deployment tools cannot rely on an eager Python module when they need a normalized tensor graph with explicit shape constraints. torch.export captures that graph as an ExportedProgram that can be inspected, serialized, and executed through PyTorch's exported-program runtime.

Example inputs define the tracing contract. The first batch dimension uses Dim(“batch”, min=1, max=8) so the loaded program accepts a batch of three even though export begins with a batch of two.

Saved export files conventionally use .pt2, but PyTorch treats their serialization format as under active development; keep producer and consumer versions aligned. torch.export.load uses pickle, so load only artifacts from a trusted source.

1. Create the model definition in export_torch_export.py.

   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)

   The ScoreModel class and its four input features are compact stand-ins for the model and tensor contract being exported.

2. Add the export input contract below the model class.

   torch.manual_seed(7)
   model = ScoreModel().eval()
   example_inputs = (torch.randn(2, 4),)
   dynamic_shapes = {"features": {0: Dim("batch", min=1, max=8)}}
    
   exported = torch.export.export(
       model,
       example_inputs,
       dynamic_shapes=dynamic_shapes,
   )

   eval() makes dropout and batch-normalization modules use inference behavior.

3. Append the artifact round-trip block below the export call.

   artifact = Path("score_model.pt2")
   torch.export.save(exported, artifact)
   loaded = torch.export.load(artifact)
    
   test_input = torch.randn(3, 4)
   eager_output = model(test_input)
   loaded_output = loaded.module()(test_input)
   torch.testing.assert_close(eager_output, loaded_output)
    
   print(f"exported type: {type(exported).__name__}")
   print(f"dynamic output shape: {tuple(loaded_output.shape)}")
   print(f"outputs match: {torch.allclose(eager_output, loaded_output)}")
   print(f"artifact saved: {artifact.exists()}")

4. Run the completed export program to verify the loaded artifact matches eager output.

   $ python export_torch_export.py
   exported type: ExportedProgram
   dynamic output shape: (3, 2)
   outputs match: True
   artifact saved: True