PyTorch 2 cheatsheet

PyTorch 2: Performance

5 common PyTorch 2 tasks.

Speed up a model

Compile the model with torch.compile for faster training and inference.

model = torch.compile(model)

Run operations in lower precision on the GPU to save memory and time.

with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
    loss = loss_fn(model(x), y)

Set the random seed for repeatable results.

torch.manual_seed(42)

Print how much GPU memory tensors use.

torch.cuda.memory_allocated() / 1e9

Release unused cached memory back to the GPU.

torch.cuda.empty_cache()