Should You Use torch.compile? A PyTorch 2.x Decision Guide

Your PyTorch 2.x script ran slower, not faster

You upgraded to PyTorch 2.0 or later, wrapped your model in torch.compile, and the first training step took noticeably longer than it did under eager execution. If your script finishes in a few minutes, that overhead can erase any benefit entirely.

The useful takeaway: torch.compile is designed for stable, long-running workloads, not short experiments. It pays an upfront compilation cost in exchange for steadier throughput later.

What happens under the hood

torch.compile rewrites eager PyTorch code into optimized kernels. TorchDynamo captures the Python bytecode, AOTAutograd generates forward and backward graphs, and Triton compiles those graphs into GPU kernels. The result is cached to disk, so subsequent runs can skip recompilation.

This feature only exists in PyTorch 2.0 and later. Check your environment first:

python -c "import torch; print(torch.__version__); print(torch.cuda.is_available())"

Run this in your normal Python environment. No special permissions are required. If the version starts with 1., torch.compile is unavailable.

Worked example: benchmark eager versus compiled

Create a small script, for example benchmark_compile.py, in your project directory. Replace YourModel and dummy_input with your actual module and data:

import time
import torch

# Replace with your model and input
model = YourModel().cuda()
example_input = torch.randn(4, 3, 224, 224).cuda()

def step(m, x):
    out = m(x)
    loss = out.sum()
    loss.backward()
    return loss

# Eager baseline
for _ in range(3):
    step(model, example_input)
torch.cuda.synchronize()

start = time.perf_counter()
for _ in range(20):
    step(model, example_input)
torch.cuda.synchronize()
eager_ms = (time.perf_counter() - start) / 20 * 1000

# Compiled variant
compiled_model = torch.compile(model, mode="default")
for _ in range(3):
    step(compiled_model, example_input)
torch.cuda.synchronize()

start = time.perf_counter()
for _ in range(20):
    step(compiled_model, example_input)
torch.cuda.synchronize()
compiled_ms = (time.perf_counter() - start) / 20 * 1000

print(f"Eager: {eager_ms:.2f} ms/step")
print(f"Compiled: {compiled_ms:.2f} ms/step")

Run it with python benchmark_compile.py. Exclude the first few compiled iterations from timing; the first run includes compilation latency and is not representative of steady-state performance. A meaningful check is whether the compiled average improves after warmup, not whether the first compiled step is fast.

To inspect whether the compiler is capturing your model or falling back to eager execution, run:

print(torch._dynamo.explain(step)(model, example_input))

Look for graph breaks in the output. Frequent breaks mean the compiler is splitting the model into many small pieces, which reduces speedup.

Trade-offs and limitations

• Cold-start latency. Compilation happens on first use. Do not judge the feature by a single short run.
• Graph breaks. Dynamic tensor shapes and Python control flow inside hot loops can force TorchDynamo to break the graph and fall back to eager execution.
• Operator coverage. Not every PyTorch operator has an optimized kernel. Unsupported operations may trigger eager fallback or errors.
• Backend maturity. CUDA support is the most mature. CPU and other backends have more limited optimization coverage.
• CUDA graphs. The reduce-overhead mode can capture CUDA graphs, which may fail with dynamic shapes or very large memory allocations.

When to enable it, and when to skip it

Use torch.compile	Avoid or test carefully
Long-running training jobs on CUDA	Short scripts that finish in seconds or minutes
Static input shapes and fixed control flow	Models with frequent dynamic shapes or Python loops
Repeated runs where disk cache can be reused	One-off notebooks and ad-hoc debugging

If compilation causes problems

torch.compile wraps your model; it does not permanently alter the model file. To disable it, simply pass the uncompiled model to your training loop. If you want to clear cached compiled artifacts, call torch._dynamo.reset() in Python, or remove the Inductor cache directory (commonly ~/.torch_inductor_cache) between runs.

Closing

Start with mode="default" on a representative training step, exclude warmup iterations from your timing, and inspect graph breaks with torch._dynamo.explain. Adopt it only if the steady-state speedup justifies the warmup cost for your specific workload.