9. Export and Deploy

This matters because a trained model is not automatically a deployable artifact. Focus on exporting a stable graph, keeping input contracts explicit, and separating training concerns from inference concerns.

import torch
from torch import nn

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = nn.Sequential(nn.Linear(8, 16), nn.ReLU(), nn.Linear(16, 2)).to(device).eval()
example = (torch.randn(4, 8, device=device),)

9.1. torch.export

exported = torch.export.export(model, example)
print(type(exported).__name__)
assert exported is not None

9.2. ONNX export

Use CPU tensors for portable ONNX export.

onnx_path = './output/exported-model.onnx'
cpu_model = model.cpu()
cpu_example = (example[0].cpu(),)
torch.onnx.export(cpu_model, cpu_example, onnx_path, input_names=['features'], output_names=['logits'], opset_version=18)
print(onnx_path)