How to compile a model in PyTorch

PyTorch models can spend extra time in Python dispatch and small operator calls after the eager model is already producing the right tensors. torch.compile wraps a function or nn.Module so PyTorch can trace tensor operations and send optimized graph regions to the compiler backend before repeated training or inference runs.

The inductor backend is the default starting point for PyTorch 2.x compilation because it balances compile overhead and runtime speed without rewriting the model class. The first compiled call can be slower while PyTorch builds the optimized path; later calls reuse cached compiled regions when shapes and control flow stay compatible.

A parity smoke test should run before the compiled callable replaces the eager path in real training or inference code. Matching output shapes and values proves the representative input still behaves the same, while separate profiling decides whether the compile overhead is worthwhile for the workload.

1. Create a compile smoke script.

   compile_smoke.py

   import torch
   from torch import nn
    
    
   class TinyClassifier(nn.Module):
       def __init__(self):
           super().__init__()
           self.layers = nn.Sequential(
               nn.Linear(8, 16),
               nn.ReLU(),
               nn.Linear(16, 3),
           )
    
       def forward(self, inputs):
           return self.layers(inputs)
    
    
   torch.manual_seed(7)
   model = TinyClassifier().eval()
   sample = torch.randn(4, 8)
   compiled_model = torch.compile(model, backend="inductor")
    
   with torch.inference_mode():
       eager_output = model(sample)
       compiled_output = compiled_model(sample)
       second_compiled_output = compiled_model(sample)
    
   max_abs_diff = (eager_output - compiled_output).abs().max().item()
    
   print(f"torch={torch.__version__}")
   print(f"input_shape={tuple(sample.shape)}")
   print(f"eager_shape={tuple(eager_output.shape)}")
   print(f"compiled_shape={tuple(compiled_output.shape)}")
   print(f"second_run_shape={tuple(second_compiled_output.shape)}")
   print(f"max_abs_diff={max_abs_diff:.8f}")
   print(f"outputs_match={torch.allclose(eager_output, compiled_output, atol=1e-6)}")

   torch.compile is available in PyTorch 2.0 and later. CPU compilation through inductor may need a C++ compiler in the active Python environment.

2. Run the smoke script.

   $ python compile_smoke.py
   torch=2.12.1+cpu
   input_shape=(4, 8)
   eager_shape=(4, 3)
   compiled_shape=(4, 3)
   second_run_shape=(4, 3)
   max_abs_diff=0.00000000
   outputs_match=True

   outputs_match=True and a zero max_abs_diff confirm the compiled callable preserved the eager model output for this input.

3. Place the compile call after model construction and checkpoint loading in the real code path.

   model = build_model()
   model.load_state_dict(torch.load("model.pt", weights_only=True))
   model.eval()
    
   compiled_model = torch.compile(model, backend="inductor")

   Keeping model and compiled_model as separate variables leaves an eager fallback available while the compiled path is being tested.
   Related: How to save and load a PyTorch model

4. Run one warm-up call with a representative input before measuring latency.

   with torch.inference_mode():
       _ = compiled_model(example_batch)

   The first compiled call includes tracing and compilation time. Measure later calls separately when comparing eager and compiled performance.

5. Use the compiled callable in the target inference call after parity checks pass.

   with torch.inference_mode():
       predictions = compiled_model(example_batch)

   Leave fullgraph at its default unless graph breaks must raise errors. Enforcing fullgraph=True can make ordinary Python control flow fail instead of allowing PyTorch to compile the capturable regions.

6. Remove the temporary smoke script after moving the pattern into the project.

   $ rm compile_smoke.py