PyTorch cheat sheet

A scannable PyTorch reference: 25 short snippets across 12 topics, each linking back to the lesson it came from.

At a glance

Quick snippets

Tensors and autograd

Autograd

x = torch.tensor(2.0, requires_grad=True)
y = x ** 3 + 2 * x

y.backward()          # compute dy/dx
x.grad                # tensor(14.)  -> 3x² + 2 at x = 2

# during evaluation you do not need the graph
with torch.no_grad():
    preds = model(inputs)

# after a step, clear the accumulated gradients
optimizer.zero_grad()

CPU, GPU and dtype

device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"

t = t.to(device)
model = model.to(device)

# every tensor in an operation must live on the same device
# RuntimeError: Expected all tensors to be on the same device

Tensors are NumPy with two extras

import torch

t = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
t.shape        # torch.Size([2, 2])
t.dtype        # torch.float32
t.device       # device(type='cpu')

torch.zeros(2, 3)
torch.ones_like(t)
torch.arange(0, 10, 2)
torch.randn(3, 3)              # standard normal

… 5 more lines in the full lesson.

Full lesson: Tensors and autograd →

The training loop

Making it converge

from torch.optim.lr_scheduler import CosineAnnealingLR

scheduler = CosineAnnealingLR(optimizer, T_max=EPOCHS)
# ... at the end of each epoch:
scheduler.step()

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)   # stabilise

Reproducibility and honest metrics

import random, numpy as np

def seed_everything(seed=42):
    random.seed(seed); np.random.seed(seed)
    torch.manual_seed(seed); torch.cuda.manual_seed_all(seed)

The canonical loop

import torch
from torch import nn

model = MyNet().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)

for epoch in range(EPOCHS):
    model.train()
    for xb, yb in train_loader:
        xb, yb = xb.to(device), yb.to(device)

… 15 more lines in the full lesson.

Full lesson: The training loop →

Modules, data and saving

Saving and loading

torch.save(model.state_dict(), "model.pt")

model = TabularNet(20, 3)
model.load_state_dict(torch.load("model.pt", map_location="cpu"))
model.eval()                       # do not forget: affects dropout and batchnorm

# resuming training keeps the optimiser too
torch.save({"model": model.state_dict(),
            "optim": optimizer.state_dict(),
            "epoch": epoch}, "checkpoint.pt")

A model is a class

import torch
from torch import nn

class TabularNet(nn.Module):
    def __init__(self, n_features, n_classes):
        super().__init__()                       # always call this first
        self.net = nn.Sequential(
            nn.Linear(n_features, 64),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(64, n_classes),
        )

… 6 more lines in the full lesson.

Dataset and DataLoader

from torch.utils.data import Dataset, DataLoader

class CSVData(Dataset):
    def __init__(self, X, y):
        self.X = torch.as_tensor(X, dtype=torch.float32)
        self.y = torch.as_tensor(y, dtype=torch.long)

    def __len__(self):
        return len(self.y)

    def __getitem__(self, i):
        return self.X[i], self.y[i]

… 3 more lines in the full lesson.

Full lesson: Modules, data and saving →

Building networks: layers, containers and initialisation

The three containers

import torch
import torch.nn as nn

class MLP(nn.Module):
    def __init__(self, sizes):
        super().__init__()
        # ModuleList, not a list: a plain Python list is invisible to .parameters()
        self.blocks = nn.ModuleList([
            nn.Linear(a, b) for a, b in zip(sizes[:-1], sizes[1:])
        ])
        self.norms = nn.ModuleList([nn.LayerNorm(b) for b in sizes[1:]])
        self.act = nn.GELU()

… 13 more lines in the full lesson.

Branches, sharing and weight tying

class TwoTower(nn.Module):
    """Shared trunk, two heads, plus a tied embedding matrix."""

    def __init__(self, vocab=1000, dim=128, n_classes=5):
        super().__init__()
        self.embed = nn.Embedding(vocab, dim)
        self.trunk = nn.Sequential(nn.Linear(dim, dim), nn.ReLU())

        # the SAME layer instance used twice: parameters are shared, not copied
        self.shared = nn.Linear(dim, dim, bias=False)

        self.head = nn.Linear(dim, n_classes)

… 13 more lines in the full lesson.

Full lesson: Building networks: layers, containers and initialisation →

Loss functions and optimisers in practice

Choosing a loss

import torch
import torch.nn as nn

logits = torch.randn(8, 5)                 # batch of 8, 5 classes
targets = torch.randint(0, 5, (8,))        # int64 class indices

# CrossEntropyLoss takes raw logits and applies log-softmax internally
ce = nn.CrossEntropyLoss()
loss = ce(logits, targets)
manual = nn.functional.cross_entropy(logits, targets)
print(torch.allclose(loss, manual))        # True

… 15 more lines in the full lesson.

Optimisers and weight decay

from torch.optim import SGD, Adam, AdamW

# never apply weight decay to norms or biases
def split_params(model, decay=0.01):
    decay_params, no_decay = [], []
    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        if param.ndim <= 1 or name.endswith(".bias"):
            no_decay.append(param)
        else:
            decay_params.append(param)

… 16 more lines in the full lesson.

Loss and optimiser failure modes

# gradient health check: run after backward, before step
def grad_report(model, top=5):
    rows = []
    for name, p in model.named_parameters():
        if p.grad is None:
            rows.append((name, "None", 0.0))
            continue
        rows.append((name, f"{p.grad.norm().item():.3e}", p.grad.abs().max().item()))
    rows.sort(key=lambda r: r[2], reverse=True)
    for name, norm, mx in rows[:top]:
        print(f"{name:40s} grad_norm={norm} max={mx:.3e}")
    total = sum(p.grad.norm().item() ** 2 for p in model.parameters() if p.grad is not None) ** 0.5

… 4 more lines in the full lesson.

Full lesson: Loss functions and optimisers in practice →

Data pipelines and augmentation with torchvision

Making the loader fast

import time

def benchmark(loader, n=30):
    it = iter(loader)
    next(it)                                  # warm up workers
    start = time.perf_counter()
    for _ in range(n):
        batch = next(it)
    elapsed = time.perf_counter() - start
    return elapsed / n, batch[0].shape

per_batch, shape = benchmark(train_loader)

… 7 more lines in the full lesson.

Full lesson: Data pipelines and augmentation with torchvision →

Transfer learning and fine-tuning

Loading a pretrained backbone

import torch
import torch.nn as nn
from torchvision import models
from torchvision.models import ResNet50_Weights

weights = ResNet50_Weights.IMAGENET1K_V2
model = models.resnet50(weights=weights)

# the transform that matches these exact weights
preprocess = weights.transforms()
print(preprocess)

… 15 more lines in the full lesson.

Progressive unfreezing and discriminative rates

def unfreeze_from(model, block_names):
    for name, child in model.named_children():
        if name in block_names:
            for param in child.parameters():
                param.requires_grad = True

def trainable_groups(model, base_lr=1e-4):
    """Give earlier layers a smaller rate than later ones."""
    groups, names = [], []
    for name, module in model.named_children():
        params = [p for p in module.parameters() if p.requires_grad]
        if params:

… 13 more lines in the full lesson.

Full lesson: Transfer learning and fine-tuning →

Learning-rate schedules and regularisation

Putting the pieces in the right order

# the order of operations in a step matters
def step(model, batch, opt, sched, scaler=None):
    x, y = batch
    opt.zero_grad(set_to_none=True)          # 1. clear, before backward
    with torch.autocast("cuda", enabled=scaler is not None):
        out = model(x)                       # 2. forward under autocast
        loss = criterion(out, y)

    if scaler is not None:
        scaler.scale(loss).backward()        # 3. scaled backward
        scaler.unscale_(opt)                 # 4. unscale before clipping
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

… 9 more lines in the full lesson.

Full lesson: Learning-rate schedules and regularisation →

Experiment tracking and reproducibility

Comparing runs honestly

# a small helper that makes comparisons explicit and cheap
import json
from pathlib import Path

def log_run(store="runs/index.jsonl", **fields):
    Path(store).parent.mkdir(parents=True, exist_ok=True)
    with open(store, "a", encoding="utf-8") as fh:
        fh.write(json.dumps(fields, default=str) + "
")

log_run(
    run="resnet18-lr3e4-bs64-seed0",

… 10 more lines in the full lesson.

Full lesson: Experiment tracking and reproducibility →

Mixed precision and GPU performance

Automatic mixed precision

import torch
from torch.amp import autocast, GradScaler

model = model.cuda()
opt = torch.optim.AdamW(model.parameters(), lr=3e-4)
scaler = GradScaler("cuda")

for x, y in loader:
    x = x.cuda(non_blocking=True)
    y = y.cuda(non_blocking=True)

    opt.zero_grad(set_to_none=True)

… 11 more lines in the full lesson.

Memory and throughput tricks

# channels_last: better tensor-core utilisation for convolutions
model = model.to(memory_format=torch.channels_last)
x = x.to(memory_format=torch.channels_last)

# pinned memory and non-blocking copies overlap transfer with compute
loader = torch.utils.data.DataLoader(
    dataset, batch_size=64, num_workers=8,
    pin_memory=True, persistent_workers=True, drop_last=True)
x = x.cuda(non_blocking=True)
y = y.cuda(non_blocking=True)

# gradient accumulation: a larger effective batch without more memory

… 14 more lines in the full lesson.

Profiling the step

from torch.profiler import profile, record_function, ProfilerActivity

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1),
    on_trace_ready=torch.profiler.tensorboard_trace_handler("logs/prof"),
    record_shapes=True,
    profile_memory=True,
) as prof:
    for step, (x, y) in enumerate(loader):
        if step >= 6:
            break

… 5 more lines in the full lesson.

Full lesson: Mixed precision and GPU performance →

Debugging PyTorch models

Shape, device and dtype errors

# a shape trace that fits in a few lines and pays for itself immediately
def trace_shapes(model, sample, name="model"):
    hooks = []

    def make_hook(module_name):
        def hook(module, inputs, output):
            in_shape = tuple(inputs[0].shape) if inputs else None
            out_shape = tuple(output.shape) if isinstance(output, torch.Tensor) else type(output).__name__
            print(f"{module_name:45s} {str(in_shape):22s} -> {out_shape}")
        return hook

    for n, m in model.named_modules():

… 11 more lines in the full lesson.

Full lesson: Debugging PyTorch models →

Exporting and deploying models

Quantisation

import torch
import torch.quantization as tq

# dynamic quantisation: weights int8, activations quantised at runtime.
# One of the few things that speeds up CPU inference substantially.
q_model = torch.quantization.quantize_dynamic(
    model.cpu().eval(), {torch.nn.Linear, torch.nn.LSTM}, dtype=torch.qint8)

def size_mb(m):
    buf = torch.save(m.state_dict(), "/tmp/m.pt")
    import os
    return os.path.getsize("/tmp/m.pt") / 1e6

… 13 more lines in the full lesson.

Packaging behind an inference API

import io
import torch
from fastapi import FastAPI, File, UploadFile
from PIL import Image
from torchvision import transforms

app = FastAPI()
model = torch.jit.load("model_traced.pt").eval()

preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),

… 16 more lines in the full lesson.

Full lesson: Exporting and deploying models →

FAQ

Related cheat sheets

AI Basics AI Agents Math for AI Machine Learning scikit-learn TensorFlow

Last refreshed 2026-09-27.