How to compile a model in PyTorch

Model execution in PyTorch is eager by default, which keeps code flexible but can leave repeated Python dispatch in the runtime path. torch.compile captures compatible tensor operations and sends them through the compiler stack while preserving the familiar nn.Module interface.

By default, torch.compile uses the TorchInductor backend and compiles lazily when the callable first receives an input. That first call includes tracing and code generation, while later calls reuse cached compiled regions until a guard such as tensor shape or Python control flow changes.

A representative input should produce the same values through both eager and compiled execution before performance measurements begin. Keeping the eager module available provides a direct reference, and torch.testing.assert_close makes a numerical mismatch stop the program instead of turning parity into a visual judgment.

1. Define the model architecture in compile_model.py.

   compile_model.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)

2. Add repeatable model state and a representative input below the class.

   compile_model.py

   torch.manual_seed(7)
   model = TinyClassifier().eval()
   sample = torch.randn(4, 8)

3. Add eager reference execution and compiled execution below the setup block.

   compile_model.py

   with torch.inference_mode():
       eager_output = model(sample)
    
   compiled_model = torch.compile(model)
    
   with torch.inference_mode():
       compiled_output = compiled_model(sample)
       repeated_output = compiled_model(sample)

   A compiled callable wraps the module state active at this point, so checkpoint loading and the training or evaluation mode belong before this line. CPU compilation can require a C++ compiler.
   Related: How to save and load a PyTorch model

4. Append numerical parity checks and result fields to the end of the file.

   compile_model.py

   torch.testing.assert_close(compiled_output, eager_output)
   torch.testing.assert_close(repeated_output, eager_output)
    
   max_abs_diff = (eager_output - compiled_output).abs().max().item()
   print(f"torch={torch.__version__}")
   print(f"input_shape={tuple(sample.shape)}")
   print(f"output_shape={tuple(compiled_output.shape)}")
   print(f"max_abs_diff={max_abs_diff:.8f}")
   print(f"compiled_matches_eager={torch.allclose(compiled_output, eager_output)}")
   print(f"repeat_matches_eager={torch.allclose(repeated_output, eager_output)}")

5. Run the completed model compilation program.

   $ python compile_model.py
   torch=2.13.0+cpu
   input_shape=(4, 8)
   output_shape=(4, 3)
   max_abs_diff=0.00000000
   compiled_matches_eager=True
   repeat_matches_eager=True

   The first compiled call builds the optimized path. Both True values plus the zero maximum difference confirm that the compiled result and its repeated call match eager execution for the representative input.