PyTorch · torchvision · torchaudio Cheat Sheet

01Install & verifyonce per env

pip install torch torchvision torchaudio★
CPU-only default wheels on PyPI.
pip install torch --index-url https://download.pytorch.org/whl/cu128
CUDA build. Pick the tag from the selector on pytorch.org.
torch.__version__
Prints e.g. 2.13.0+cu128 — build tag included.
torch.cuda.is_available()
Quick sanity check that the GPU build actually sees a device.
python -m torch.utils.collect_env
Full environment dump — paste this into bug reports.
torch / torchvision / torchaudio versionspin
The domain packages link against one exact torch build. Install all three together or you get import-time symbol errors.
torch.cuda.device_count() · get_device_name(0)
How many GPUs, and which ones.
torch.backends.mps.is_available()
Apple Silicon. 2.13 brings FlexAttention to MPS.
pip install torchcodecrequired
Separate package. Every audio and video decode path now goes through it.

02The import mapwhat lives where

import torch★
Tensors, math, devices, save/load, compile.
import torch.nn as nn★
Layer classes that hold parameters.
import torch.nn.functional as F
The same ops as stateless functions.
import torch.optim as optim
SGD, Adam, AdamW + LR schedulers.
from torch.utils.data import Dataset, DataLoader★
The two halves of every input pipeline.
import torch.distributed as dist
Multi-GPU / multi-node collectives.
from torchvision.transforms import v2
Note the v2 — see card V2.
import torchaudio
Waveforms, spectrograms, ASR pipelines.
from torch import Tensor
For type hints: def f(x: Tensor) -> Tensor:
import torchvision.transforms.v2.functional as TF
Stateless versions when you need to apply the same op to two things.
from torch.utils.tensorboard import SummaryWriter
Built in. w.add_scalar("loss/train", v, step).

03Create tensorsthe only data type

torch.tensor(data)★
From a list or ndarray. Copies, and infers the dtype.
torch.zeros(2, 3) · ones · empty · full★
Shape given as loose ints or a tuple.
torch.randn(2, 3) · rand · randint★
randn is N(0,1); rand is uniform [0,1).
torch.arange(0, 10, 2) · linspace(0, 1, 5)
Step-based vs count-based ranges.
torch.eye(3) · torch.from_numpy(arr)
from_numpy shares memory with the ndarray.
torch.zeros_like(x) · randn_like(x)
Copies shape, dtype and device from x.
x.clone()
Deep copy that stays in the autograd graph. Add .detach() to leave it.
torch.manual_seed(0)
Seeds all devices. Do it before you build the model.
torch.zeros(2, 3, device=dev, dtype=torch.float32)★
Build it on the device. Creating on CPU then .to(dev) allocates twice and copies.
torch.as_tensor(arr)
Shares memory with the numpy array where it can. torch.tensor always copies.
torch.randperm(n) · multinomial(w, k)
Random permutation; weighted sampling without replacement.
g = torch.Generator().manual_seed(0)
A local RNG. Pass generator=g to DataLoader and samplers so seeding survives workers.

04Attributes & dtypesread before you debug

x.shape · x.ndim · x.numel()★
Dimensions, rank, total element count.
x.dtype · x.device · x.requires_grad★
The three things wrong when a call fails.
x.to(torch.float32) · x.float() · x.long()
Cast. .to() is a no-op if it already matches.
float32 · float16 · bfloat16
Default is float32. bfloat16 keeps float32's exponent range — safer for training than float16.
int64 · bool · uint8
Class labels must be int64; images arrive as uint8.
x.item()
One-element tensor → Python number. Forces a device sync.
x.tolist() · x.numpy()
.numpy() needs CPU and no grad: x.detach().cpu().numpy().
x.is_contiguous()
False after permute/transpose — and then .view() raises.
x.element_size() * x.numel()
Bytes held by the tensor — the fastest memory estimate you have.
torch.finfo(torch.float16).max
Dtype limits. Handy when hunting an inf in half precision.
x.stride()
Step per dimension. A permute changes stride, not memory.
int tensor + float tensorquiet
Dtypes promote silently. An accidental float64 anywhere will quietly halve your throughput.

05Shape surgerywhere most bugs live

x.view(2, -1)★
Free re-interpretation. Needs contiguous memory; -1 is inferred.
x.reshape(2, -1)★
Same result, but copies if it has to. Use this when unsure.
x.permute(2, 0, 1)
Reorder all dims — e.g. HWC → CHW.
x.transpose(0, 1) · x.T
Swap exactly two dims.
x.unsqueeze(0) · x.squeeze()
Add / drop size-1 dims. unsqueeze(0) is how you fake a batch.
x.flatten(start_dim=1)
Collapse everything after the batch dim — the CNN→Linear join.
torch.cat([a, b], dim=0)★
Joins along an existing dim. Rank unchanged.
torch.stack([a, b], dim=0)★
Creates a new dim. Rank +1. This is the pair people mix up.
x.expand(3, -1) · x.repeat(3, 1)
expand is a free view; repeat allocates.
x.contiguous()
The fix for "view size is not compatible…".
torch.split(x, 2, dim=0) · chunk(x, 3)
split takes a size, chunk takes a count.
x.movedim(1, -1)
Move one axis and leave the rest alone — clearer than spelling out a full permute.
nn.Flatten(1) · nn.Unflatten(1, (C, H, W))
The layer forms, for use inside nn.Sequential.

06Broadcastingthe silent bug factory

shapes align from the right★
Trailing dims must be equal, or one of them 1, or absent. That one rule explains every broadcast.
(3,1) + (1,4) → (3,4)★
Size-1 dims stretch. Nothing is copied until it has to be.
(8,3,32,32) * (3,1,1)
Per-channel scale on a batch of images. The classic normalize shape.
torch.broadcast_shapes(a.shape, b.shape)
Ask before you multiply. Raises if they are incompatible.
x.unsqueeze(1) * y
Be explicit about where the new axis goes rather than trusting alignment.
x.sum(dim=1, keepdim=True)
Keeps rank so the result still broadcasts back against the original.
pred (N,1) vs target (N,)★famous
The expensive one. These broadcast to (N,N), so your loss silently averages N² pairs and trains on nonsense. Nothing raises.

07Index, slice, reduceNumPy rules apply

x[0] · x[:, 1] · x[..., -1]★
Basic slicing returns a view, not a copy.
x[x > 0]
Boolean mask → flat 1-D tensor of the matches.
torch.where(cond, a, b)
Elementwise pick — differentiable, unlike a Python if.
x.sum(dim=1) · mean · std · prod★
Named dim disappears unless keepdim=True.
x.max(dim=1)
Returns (values, indices) — a named tuple, not a bare tensor.
x.argmax(dim=1)★
Logits → predicted class. The classification one-liner.
torch.topk(x, k=5) · sort · unique
Top-5 accuracy, ranking, vocabulary building.
x.gather(1, idx) · x.scatter_(1, idx, v)
Pick / write per-row by index. Trailing _ means in-place.
torch.nonzero(x, as_tuple=True)
Indices where the condition holds; the tuple form indexes directly.
x.masked_fill_(mask, float("-inf"))★
How attention masks are applied — −inf before softmax becomes exactly 0 after.
x.any(dim) · x.all(dim) · x.count_nonzero()
Boolean reductions.
torch.cumsum(x, dim) · roll · flip
Running totals, circular shift, reverse along an axis.

08Math & linear algebrathe actual compute

a + b · a * b · a ** 2★
Elementwise, with NumPy broadcasting.
a @ b · torch.matmul(a, b)★
Matrix multiply; batches over leading dims.
torch.einsum("bij,bjk->bik", a, b)
Names every axis. Worth learning for attention code.
torch.bmm(a, b) · mv · dot · outer
Explicit batch-matrix, matrix-vector, vector forms.
torch.exp · log · sqrt · abs · clamp
clamp(min=, max=) is the standard gradient guard.
F.softmax(x, dim=-1) · F.log_softmax
Always name the dim. log_softmax is the numerically stable one.
torch.linalg.norm · inv · svd · solve
The modern namespace — the bare torch.svd spellings are legacy.
a.add_(b)in-place
Trailing _ mutates. Autograd will raise if the value was needed for backward.
torch.allclose(a, b, atol=1e-6)★
Compare floats this way in tests. == on floats will eventually embarrass you.
torch.isnan(x).any() · nan_to_num(x)★
First thing to run when a loss goes to NaN. Check the inputs before blaming the model.
F.normalize(x, dim=-1)
Unit-length rows. Then a @ b.T is cosine similarity.
torch.cdist(a, b) · corrcoef
Pairwise distances; correlation matrix.

09Devices & precisioncuda · mps · xpu · cpu

torch.accelerator.current_accelerator()modern
Device-agnostic: returns cuda / mps / xpu without an if-chain.
torch.accelerator.is_available()
Guard the call above; falls back to "cpu".
x = x.to(dev) · model.to(dev)★
Tensors return a copy; modules move in place. Don't write model = model.to() expecting otherwise.
with torch.autocast(dev, dtype=torch.bfloat16):★
Mixed precision. Wrap the forward + loss only.
scaler = torch.GradScaler()
Needed for float16, not for bfloat16.
torch.set_float32_matmul_precision("high")
Enables TF32 matmuls on Ampere+. Near-free speedup.
x.to(dev, non_blocking=True)
Async copy — only helps with pin_memory=True loaders.
torch.cuda.empty_cache()
Returns cached blocks to the driver. Does not fix a real leak.
RuntimeError: expected all tensors on the same device
The single most common runtime error. Move both the batch and the model.
torch.cuda.synchronize()★
CUDA calls are asynchronous. Without this your timing measures how fast Python queued work, not how fast it ran.
torch.cuda.max_memory_allocated() / 1e9
Peak GB. Pair with reset_peak_memory_stats() to measure one region.

10Autogradthe mechanism

x = torch.tensor(2., requires_grad=True)★
Opt a leaf tensor into tracking. Model parameters do this for you.
loss.backward()★
Accumulates d(loss)/d(param) into every .grad. Loss must be a scalar.
x.grad
None until the first backward. Only leaves keep it.
with torch.no_grad():★
Stop recording — for evaluation and manual parameter edits.
with torch.inference_mode():★
Stricter and faster than no_grad. Prefer it for pure inference.
x.detach()
Same data, cut out of the graph. The way to stop a gradient mid-network.
torch.autograd.grad(y, x)
Returns grads instead of writing .grad — for higher-order tricks.
loss.backward(retain_graph=True)rare
The graph is freed after backward. Only set this if you truly need a second pass — otherwise it silently leaks memory.
torch.autograd.set_detect_anomaly(True)★
Points at the forward op that produced a NaN gradient. Slow — switch on only while debugging.
torch.func.grad(f)(x) · vmap · jacrev
Functional autograd: per-sample gradients, jacobians, batching without a loop.
total = nn.utils.clip_grad_norm_(params, 1.0)
It returns the pre-clip norm — log it. A spike there precedes most loss explosions.
t.register_hook(lambda g: print(g.norm()))
Inspect a gradient partway through the graph.

11nn.Moduleevery model is one

class Net(nn.Module):★
Subclass, then define __init__ and forward. Never call forward() yourself — call model(x).
super().__init__()
First line of __init__. Skip it and no parameter registers.
nn.Sequential(layer, layer, ...)★
A model without a class, when the graph is a straight line.
nn.ModuleList([...]) · nn.ModuleDict
Use these, not a plain Python list — a plain list hides the parameters.
nn.Parameter(torch.randn(3))
A tensor attribute that counts as a trainable weight.
self.register_buffer("mask", t)
Moves with .to() and saves in the state dict, but is not trained.
model.parameters() · model.named_parameters()★
What you hand the optimizer.
sum(p.numel() for p in model.parameters())
Parameter count in one line.
model.apply(init_fn) · nn.init.kaiming_normal_
Custom weight init, applied recursively.
for p in model.parameters(): p.requires_grad_(False)★
Freeze a backbone. Then unfreeze the head and hand only those params to the optimizer.
nn.Identity()
Replace a head with a pass-through instead of deleting it — shapes downstream stay intact.
model.register_forward_hook(fn)
Grab intermediate activations without editing forward.
model.named_modules() · model.children()
Walk the tree. named_modules recurses, children is one level.

12Layersnn.* building blocks

nn.Linear(in_f, out_f)★
Acts on the last dim only; leading dims pass through.
nn.Conv2d(in_ch, out_ch, k, stride, padding)★
Expects (N, C, H, W). padding="same" preserves size at stride 1.
nn.MaxPool2d(2) · nn.AdaptiveAvgPool2d(1)
The adaptive one fixes the output size for any input — how CNNs reach a Linear head.
nn.BatchNorm2d(ch) · nn.LayerNorm(shape)★
BatchNorm across the batch (CNNs); LayerNorm within a sample (transformers).
nn.Dropout(p=0.5)
Active in train(), a no-op in eval().
nn.Embedding(num_emb, dim)
Integer ids → vectors. A lookup, not a matmul.
nn.MultiheadAttention(dim, heads)
Pass batch_first=True unless you like (L, N, E).
nn.TransformerEncoderLayer · LSTM · GRU
RNNs return (output, hidden), not a bare tensor.
F.scaled_dot_product_attention(q, k, v)fast
Fused attention kernel. torch.nn.attention.flex_attention handles custom masks — and as of 2.13 runs on Apple Silicon too.
out = (in + 2*pad − k) // stride + 1★
The conv size formula. Memorise it and you stop guessing why a Linear layer rejects your flattened feature map.
nn.GroupNorm(g, ch) · nn.RMSNorm(dim)modern
GroupNorm when batches are tiny; RMSNorm is the norm of choice in modern LLMs.
nn.ConvTranspose2d · nn.Upsample(scale_factor=2)
Going back up in resolution. Upsample+Conv avoids checkerboard artefacts.
nn.Conv1d · Conv3d · padding="same"
1D for sequences and audio, 3D for volumes.

13Losses & activationspick by task

nn.CrossEntropyLoss()★
Multi-class. Takes raw logits and int64 targets — softmax is built in.
nn.BCEWithLogitsLoss()★
Binary / multi-label. Also takes logits, and is stabler than Sigmoid + BCELoss.
nn.MSELoss() · nn.L1Loss() · SmoothL1Loss★
Regression. Smooth L1 (Huber) resists outliers.
nn.NLLLoss() · KLDivLoss · CTCLoss
NLLLoss expects log_softmax output; CTCLoss is the speech/OCR one.
loss_fn(pred, y, reduction="none")
Per-sample losses — needed for weighting and focal variants.
nn.ReLU() · GELU · SiLU · LeakyReLU★
GELU/SiLU are the transformer defaults; ReLU still fine for CNNs.
nn.Sigmoid() · Tanh · Softmax(dim=-1)
For reading probabilities at inference — not before a logits-based loss.
nn.Softmax() → nn.CrossEntropyLoss()wrong
Double softmax. Trains, converges badly, and nothing warns you.
nn.CrossEntropyLoss(label_smoothing=0.1)★
Nearly free accuracy on most classification tasks. Stops the model getting overconfident.
nn.CrossEntropyLoss(weight=w, ignore_index=-100)★
Class imbalance, and skipping pad tokens. -100 is the conventional pad label.
nn.BCEWithLogitsLoss(pos_weight=w)
Imbalance for binary and multi-label problems.
nn.LinearCrossEntropyLoss()2.13
Fuses the final projection into the loss. Up to 4× less peak memory on large-vocabulary LM training.
F.one_hot(y, num_classes=n)
Rarely needed — CrossEntropyLoss wants class indices, not one-hot.

14Optimizers & schedulestorch.optim

opt = optim.AdamW(model.parameters(), lr=1e-3)★
The sane default. Decouples weight decay properly — prefer it to Adam.
optim.SGD(params, lr=0.1, momentum=0.9)
Still the best final accuracy for vision, with a schedule.
opt.zero_grad(set_to_none=True)★
Default since 2.0 — frees memory as well as clearing.
opt.step()★
Applies whatever is in .grad right now.
opt = optim.AdamW([{"params": a, "lr": 1e-4}, ...])
Param groups: different LR for backbone vs head.
nn.utils.clip_grad_norm_(params, 1.0)
Between backward() and step(). Standard for RNNs and LLMs.
optim.lr_scheduler.CosineAnnealingLR(opt, T_max)
The common modern decay.
OneCycleLR · StepLR · ReduceLROnPlateau
ReduceLROnPlateau is the only one whose step() takes a metric.
sched.step()order
Call it after opt.step(), once per epoch (or per batch for OneCycle) — never before.
optim.AdamW(params, lr=3e-4, weight_decay=0.01)★
AdamW decouples decay from the gradient — this is why you pick it over Adam.
decay / no-decay param groups
Convention: no weight decay on biases and norm parameters. Split them into two groups.
optim.AdamW(params, ..., fused=True)fast
One fused CUDA kernel for the whole step. Free speed on GPU.
SequentialLR(opt, [LinearLR(...), CosineAnnealingLR(...)], milestones=[warm])★
Warmup then decay, the standard modern schedule.
sched.get_last_lr()[0]
Log it every epoch. A schedule you cannot see is a schedule you cannot debug.

15Dataset & DataLoadertorch.utils.data

class DS(Dataset): __len__, __getitem__★
Those two methods are the entire contract. Return one sample, not a batch.
TensorDataset(X, y)
Wraps tensors you already have in memory.
DataLoader(ds, batch_size=32, shuffle=True)★
Shuffle the train split; never the validation split.
num_workers=4, pin_memory=True★
The two flags that usually fix a GPU-starved loop.
persistent_workers=True, prefetch_factor=2
Avoids re-spawning workers every epoch.
drop_last=True
Drops the ragged final batch — matters for BatchNorm.
collate_fn=my_fn
Custom batching: padding variable-length text or audio.
random_split(ds, [0.8, 0.2]) · Subset · ConcatDataset
Splitting and combining without touching disk.
WeightedRandomSampler(w, n)
Class imbalance. Mutually exclusive with shuffle=True.
DataLoader(ds, sampler=s, shuffle=False)
sampler and shuffle are mutually exclusive — passing both raises.
class S(IterableDataset): __iter__
For streams with no length — shard by get_worker_info() or every worker yields the same data.
DataLoader(..., generator=g, worker_init_fn=fn)
What actually makes a shuffled, multi-worker loader reproducible.

16The training loopno framework required

model.train()★
Turns Dropout on and BatchNorm into batch-statistics mode. It does not enable gradients.
for x, y in loader:
The DataLoader is a plain iterable; one pass = one epoch.
pred = model(x); loss = loss_fn(pred, y)★
Forward. The graph is built here, fresh, every iteration.
opt.zero_grad(); loss.backward(); opt.step()★
The three lines. Order matters; grads accumulate if you skip the first.
running += loss.item() * x.size(0)
Log the float, not the tensor — keeping the tensor keeps its whole graph alive.
loss = loss / accum_steps
Gradient accumulation: skip step() for N-1 iterations to fake a bigger batch.
model.eval()
The other half of the pair. Forgetting it is why validation accuracy looks wrong.
with torch.inference_mode(): ...
Wrap the whole eval pass. Roughly halves memory.
acc = (logits.argmax(1) == y).float().mean()
Accuracy without a metrics library.
with torch.autocast(dev, dtype=torch.bfloat16):★
Mixed precision, three lines. Forward inside the context; backward outside it.
scaler.scale(loss).backward(); scaler.step(opt); scaler.update()
The float16 path only. bfloat16 has the range to skip the scaler entirely.
@torch.inference_mode() def evaluate(...)★
Decorate the whole eval function rather than remembering the with block.
if val < best: best = val; torch.save(...)
Track the best, not the last. The final epoch is rarely the best one.

17Save, load, checkpointstate dicts

torch.save(model.state_dict(), "m.pt")★
Save the weights, not the object. A pickled model breaks when your code moves.
model.load_state_dict(torch.load("m.pt"))★
Build the architecture first, then pour the weights in.
torch.load(p, weights_only=True)default
The default since 2.6. It blocks arbitrary code execution — pass False only for files you trust.
torch.load(p, map_location="cpu")★
Loads a GPU checkpoint on a CPU-only machine.
torch.save({"model": ..., "opt": ..., "epoch": e}, p)
A resumable checkpoint needs the optimizer state too.
model.load_state_dict(sd, strict=False)
Returns the missing / unexpected keys. Useful for partial transfer.
keys prefixed "module."
Saved from DataParallel/DDP. Strip the prefix or save model.module.state_dict().
keys prefixed "_orig_mod."★compile
The modern twin of the module. bug. torch.compile wraps your model, so the saved keys gain a prefix. Save model._orig_mod.state_dict(), or strip it on load.
torch.load(p, mmap=True)
Memory-maps the file. Large checkpoints load without doubling RAM.
ckpt = {"model", "opt", "sched", "scaler", "epoch"}★
A resumable checkpoint is all five. Saving only weights means restarting the schedule from zero.

18torch.compile & speedthe 2.x headline

model = torch.compile(model)★
Traces to a graph and fuses kernels. Usually a one-line speedup; keep it after .to(dev).
torch.compile(model, mode="max-autotune")
Slower to compile, faster to run. Also "reduce-overhead" for small models.
torch.compile(model, dynamic=True)
Avoids recompiling for every new sequence length.
TORCH_LOGS="recompiles" python train.py
Diagnose why compilation keeps re-triggering.
channels_last memory format
model.to(memory_format=torch.channels_last) — real gains for CNNs on tensor cores.
with torch.profiler.profile(...) as p:
Per-op timings, exportable to TensorBoard or Chrome trace.
graph breaks
Prints, .item(), and data-dependent Python control flow all split the graph and cost most of the win.
torch.compile(model, fullgraph=True)★
Fail loudly on a graph break instead of quietly falling back to eager.
torch.backends.cudnn.benchmark = True
Autotunes conv algorithms. A win for fixed input shapes, a loss for varying ones.
start = torch.cuda.Event(enable_timing=True)
The correct way to time GPU work. time.time() around async kernels measures nothing.

19Memory & OOMwhen CUDA runs out

running_loss += loss★classic
Reach for this first. Accumulating the tensor keeps its whole autograd graph alive, so memory grows every step. Use loss.item() or .detach().
loss = loss / accum; if (i+1) % accum == 0: opt.step()★
Gradient accumulation. Large effective batch on a small card, for the price of time.
torch.utils.checkpoint.checkpoint(block, x)★
Discard activations and recompute them in the backward pass. Roughly 30% slower, dramatically smaller.
with torch.autocast(dev, dtype=torch.bfloat16):
Halves activation memory and usually speeds things up.
opt.zero_grad(set_to_none=True)
Frees the gradient tensors instead of filling them with zeros. Default since 2.0.
del logits, loss; torch.cuda.empty_cache()
Returns cached blocks to the driver. Helps fragmentation, not a true leak.
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
Environment-variable fix for “tried to allocate” errors when plenty of memory is free.
torch.cuda.memory_summary()
Where it all went, by allocation site.

20Scale out & shipmulti-GPU · export

torchrun --nproc_per_node=8 train.py★
The launcher. It sets RANK, LOCAL_RANK, WORLD_SIZE for you.
model = DistributedDataParallel(model)★
One process per GPU. Always prefer it to the older DataParallel.
DistributedSampler(ds)
Required with DDP, and you must call set_epoch(e) or every epoch shuffles identically.
FSDP2 · fully_shard
Shards parameters, grads and optimizer state — for models too big for one card.
ep = torch.export.export(model, (x,))modern
The current capture path, ahead of TorchScript.
torch.onnx.export(model, (x,), "m.onnx", dynamo=True)
ONNX for cross-runtime serving.
torch.jit.script · torch.jit.tracelegacy
TorchScript still works but is no longer the recommended path — new code should use torch.export.
dist.init_process_group("nccl") ... destroy_process_group()★
Open and close it. Skipping the teardown is how jobs hang at exit.
dist.all_reduce(t, op=dist.ReduceOp.SUM)
Aggregate metrics across ranks — otherwise every rank logs its own shard.
if dist.get_rank() == 0: torch.save(...)★
Guard saving and logging, or eight ranks race to write one file.
torchcomms2.13
New distributed backend: better fault tolerance and debuggability on large clusters.

21Debug & reproducebefore you blame the model

print(x.shape, x.dtype, x.device)★
The three questions. Almost every runtime error in PyTorch is one of these three disagreeing.
overfit a single batch first★
If the model cannot drive one batch to near-zero loss, the bug is in the model or the loss — not the data or the schedule.
torch.manual_seed(s); random.seed(s); np.random.seed(s)★
Three separate RNGs. Seeding one of them is the usual mistake.
torch.use_deterministic_algorithms(True)
Raises on any op with no deterministic implementation, rather than silently varying.
CUBLAS_WORKSPACE_CONFIG=:4096:8
Required alongside the flag above, or CUDA matmuls still wander.
assert not torch.isnan(loss)
Fail on the step that broke, not three epochs later when everything is NaN.
torch.autograd.set_detect_anomaly(True)
Traces a NaN gradient back to the forward op that made it.
model.eval() forgotten★
Dropout stays on and BatchNorm keeps updating its statistics. Validation accuracy looks noisy and low for no visible reason.

V1What's insidesix namespaces

torchvision.transforms.v2★
Augmentation for images, boxes, masks and video.
torchvision.models★
Pretrained classification, detection, segmentation, video backbones.
torchvision.datasets
CIFAR, MNIST, ImageNet, COCO, plus ImageFolder.
torchvision.ops
Boxes, NMS, RoI pooling, focal loss.
torchvision.io
Image decode / encode. Video was removed here.
torchvision.tv_tensors
Tensor subclasses that tell v2 transforms what they're holding.
torchvision.utils
Grids, saving, and drawing boxes / masks / keypoints.
torchvision.models.feature_extraction
Pull named intermediate layers out of any model as a new module.
torchvision.transforms.v2.functional
Stateless ops when you must apply the identical transform to image and mask.

V2transforms v2 · the pipelinethe canonical order

from torchvision.transforms import v2★
Import v2, not transforms. Same names, faster, and it handles boxes and masks.
v2.Compose([...])★
Chain them. Order is not cosmetic — see the four lines below.
v2.ToImage()
1. PIL / ndarray → an Image tv_tensor. Replaces the old ToTensor.
v2.RandomResizedCrop(224, antialias=True)
2. Geometry and augmentation, while still cheap uint8.
v2.ToDtype(torch.float32, scale=True)
3. scale=True maps [0,255] → [0,1]. Omit it and your images stay in the hundreds.
v2.Normalize(mean=[...], std=[...])
4. Last. ImageNet stats: mean .485 .456 .406, std .229 .224 .225.
v2.Lambda(fn) · v2.RandomApply([...], p=0.3)
Escape hatches for anything custom.
v2.ToTensor()retired
Did convert and rescale in one opaque step. Use ToImage() + ToDtype(..., scale=True).
train pipeline ≠ eval pipeline★
Build two. Random crops and flips belong to training only; eval gets Resize + CenterCrop and nothing stochastic.
geometry → ToImage → ToDtype → Normalize★
The order that works. Augment while cheap uint8, convert once, normalize last.
antialias=True
The default in v2, and the reason v2 resizing matches PIL where v1 did not.

V3Transforms worth knowinggeometry · colour

v2.Resize(256) · v2.CenterCrop(224)★
The standard eval pair. A single int resizes the shorter side.
v2.RandomHorizontalFlip(p=0.5)★
Cheapest real augmentation there is. Wrong for text and digits.
v2.RandomRotation · RandomAffine · RandomPerspective
Geometric jitter; each takes a fill for the new corners.
v2.ColorJitter(brightness=.4, contrast=.4)
Photometric jitter. Skip it when colour is the signal.
v2.RandAugment() · TrivialAugmentWide() · AutoAugment()★
Learned policies. TrivialAugmentWide is parameter-free and a strong default.
v2.RandomErasing(p=0.25)
Cutout. Must come after ToDtype.
v2.CutMix(num_classes=n) · v2.MixUp(...)
Batch-level: apply after the DataLoader or inside collate_fn, not per sample.
v2.functional.rotate(img, 30)
Functional forms, for applying identical params to an image and its mask.
v2.RandomCrop(32, padding=4)★
The CIFAR standard, paired with a horizontal flip.
v2.GaussianBlur(k) · v2.RandomGrayscale(p) · v2.RandomSolarize
The self-supervised augmentation family (SimCLR, BYOL).
v2.RandomChoice([...]) · v2.RandomOrder([...])
Pick one at random, or shuffle the order each call.
v2.Resize(s, interpolation=InterpolationMode.LANCZOS)0.28
Lanczos resampling.

V4TVTensorsdetection & segmentation

from torchvision import tv_tensors
Tensor subclasses. Zero copy — they only add a label so v2 can dispatch.
tv_tensors.Image(t) · tv_tensors.Video(t)
Plain tensors are treated as images anyway; wrap for clarity.
tv_tensors.BoundingBoxes(b, format="XYXY", canvas_size=hw)★
This is what makes a flip move the boxes too.
tv_tensors.Mask(m)
Segmentation masks get nearest-neighbour interpolation automatically.
img, target = transform(img, target)★
v2 transforms take arbitrary structures — dicts, tuples, lists — and transform them together.
v2.SanitizeBoundingBoxes()
Drops boxes a crop pushed off-canvas, and the matching labels and masks. Put it last.
v2.ClampBoundingBoxes() · v2.ConvertBoundingBoxFormat(...)
Clip to the image, or switch XYXY / XYWH / CXCYWH.
tv_tensors.wrap(t, like=orig)
Re-wrap a plain tensor after an op that dropped the subclass.
plain tensors pass through untouched★
v2 transforms only apply geometry to things it recognises. That is exactly why the wrapper types exist.
tv_tensors.KeyPoints(pts, canvas_size=hw)
Landmarks and pose points, transformed in step with the image.

V5Datasetsbuilt in

datasets.ImageFolder(root, transform=tf)★
One folder per class. The fastest way to train on your own images.
datasets.CIFAR10(root, train=True, download=True)★
Same signature for MNIST, FashionMNIST, CIFAR100.
ds.classes · ds.class_to_idx
Label names and the mapping — needed to read predictions back.
datasets.CocoDetection · VOCSegmentation · ImageNet
Detection and segmentation sets; several need a manual download.
wrap_dataset_for_transforms_v2(ds)
Makes a built-in detection dataset emit tv_tensor targets.
datasets.DatasetFolder(root, loader, extensions)
ImageFolder generalised to any file type.
root/class_a/*.jpg · root/class_b/*.jpg★
The layout ImageFolder expects. Directory names become the labels, sorted alphabetically.
datasets.MNIST · FashionMNIST · Flowers102 · Food101
Small ones worth knowing for a quick sanity run.
video datasets now decode through TorchCodec0.26
Migrated transparently when video I/O left torchvision. Your code did not change.

V6Models & weightsthe enum API

from torchvision.models import resnet50, ResNet50_Weights★
Every architecture ships a matching *_Weights enum.
resnet50(weights=ResNet50_Weights.DEFAULT)★
Best available checkpoint. Note it can change between releases — pin the explicit tag for reproducibility.
resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)
Explicit and stable. Strings work too: weights="IMAGENET1K_V2".
resnet50(weights=None)
Random init — the architecture only.
preprocess = weights.transforms()★
The exact preprocessing that checkpoint was trained with. Guessing it costs real accuracy.
weights.meta["categories"]
Class-index → human label, shipped with the weights.
models.list_models() · models.get_model("resnet50", weights=...)
Look up by string — handy for config-driven training.
model.fc = nn.Linear(model.fc.in_features, n)★
Transfer learning. The head is .fc on ResNet, .classifier on VGG/EfficientNet, .heads on ViT.
resnet50(pretrained=True)removed
Gone since 0.15. Every tutorial written before 2023 uses it — translate to weights=.
create_feature_extractor(m, {"layer4": "feat"})★
Returns a module giving you named intermediate outputs. Far cleaner than hooks for embeddings.
weights.meta["num_params"] · weights.meta["_metrics"]
Parameter count and reported accuracy, carried on the enum itself.
models.detection · models.segmentation · models.video
The other model families — same weights= API throughout.

V7ops & utilsdetection plumbing

ops.nms(boxes, scores, iou_threshold)★
Returns kept indices, sorted by score. Boxes must be XYXY.
ops.batched_nms(boxes, scores, idxs, thr)
Per-class NMS — suppression never crosses categories.
ops.box_iou(a, b) · ops.generalized_box_iou
Pairwise N×M IoU matrix.
ops.box_convert(b, "xywh", "xyxy") · ops.box_area
XYXY is the format every torchvision utility expects.
ops.roi_align · ops.MultiScaleRoIAlign
Crop features per box; the FPN version picks the level for you.
ops.sigmoid_focal_loss(logits, targets)
RetinaNet's answer to dense-detection class imbalance.
ops.masks_to_boxes · ops.remove_small_boxes
Segmentation → detection targets, and degenerate-box cleanup.
utils.make_grid(batch, nrow=8) · utils.save_image★
Eyeball a batch in one line.
utils.draw_bounding_boxes · draw_segmentation_masks★
Annotates a uint8 CHW tensor. Passing float images is the usual mistake.
ops.Conv2dNormActivation · ops.StochasticDepth
Building blocks lifted out of the torchvision model implementations.
ops.complete_box_iou_loss · distance_box_iou_loss
CIoU and DIoU regression losses for detection heads.
ops.box_iou(a, b) with rotated boxes0.28
Rotated-box support.

V8Image I/Oand where video went

io.decode_image("cat.jpg")★
File or byte tensor → uint8 CHW tensor. Format-agnostic.
io.read_image(p, mode=io.ImageReadMode.RGB)
Force 3 channels — otherwise PNGs may arrive with an alpha band.
io.write_jpeg(t, p) · io.write_png(t, p)
Encoders stay in torchvision.
io.decode_jpeg(data, device="cuda")
nvJPEG GPU decode — worth it for large-batch pipelines.
io.read_video · io.write_video · io.VideoReaderremoved
Deprecated in 0.22, removed in 0.26. Use torchcodec.decoders.VideoDecoder. The video datasets were migrated for you.
data = io.read_file(p); img = io.decode_image(data)
Split read from decode when the bytes come from S3 or a database rather than a path.
io.encode_jpeg(t, quality=90)
Back to bytes without touching the filesystem.
TF.to_pil_image(t) · TF.pil_to_tensor(img)
The PIL bridge, for when a library hands you the wrong one.

V9Detection & segmentationthe other two heads

m = fasterrcnn_resnet50_fpn_v2(weights=W.DEFAULT); m.eval()★
Same weights enum as classification. The _v2 variants are the better-trained ones.
preds = m([img1, img2])★
List in, list out. Detection models take a list of images of differing sizes, not a stacked batch.
preds[0]["boxes" | "labels" | "scores"]
Boxes are XYXY in pixels, already sorted by score.
keep = preds[0]["scores"] > 0.5
Nothing is thresholded for you — the raw output is full of low-confidence boxes.
m.train(); losses = m(imgs, targets); sum(losses.values())★
In train mode the same model returns a dict of losses instead of predictions. Sum them yourself.
maskrcnn_resnet50_fpn_v2 → preds[0]["masks"]
Instance masks, shape (N,1,H,W), soft — threshold at 0.5.
deeplabv3_resnet50(...)["out"].argmax(1)
Semantic segmentation returns a dict; take "out" and argmax over classes.
utils.draw_bounding_boxes(img_uint8, boxes, labels)★
Wants uint8 input. Feeding it your normalized float tensor gives a black rectangle.

A1What survivedthe 2.9 cut

torchaudio.transforms★
Feature extraction as nn.Modules. The heart of the package.
torchaudio.functional★
The same operations, stateless.
torchaudio.pipelines★
Bundled model + weights + the exact preprocessing.
torchaudio.models · .models.decoder
Wav2Vec2, HuBERT, Conformer, Emformer RNN-T + CTC decoders.
torchaudio.datasets · .compliance.kaldi
LibriSpeech and friends; Kaldi-compatible fbank / MFCC.
torchaudio.io · torio · StreamReader · apply_effects_fileremoved
Deprecated in 2.8, deleted in 2.9. There is no shim.
lfilter · RNNTLoss · forced_align · overdrive
Slated for removal, then kept after user pushback. Safe to use.
lfilter · overdrive now pure-PythonC++ dropped
Kept after community pushback, but reimplemented without C++. Correct, and slower than they were.
goal: a Python-only library
The whole refactor exists to delete the C++ extension and shrink the dependency surface.

A2Load & savenow via TorchCodec

waveform, sr = torchaudio.load("a.wav")★
Still works — but it is now an alias for load_with_torchcodec(), so you need torchcodec installed.
torchaudio.save("out.wav", waveform, sr)
Likewise an alias for save_with_torchcodec().
from torchcodec.decoders import AudioDecoderpreferred
The native path the team recommends porting to.
pip install torchcodec
Separate package, not pulled in by torchaudio.
waveform.shape → (channel, time)★
Channel-first, float32 in [-1, 1]. librosa hands you (time,) — that mismatch causes silent shape bugs.
waveform.mean(dim=0, keepdim=True)
Downmix stereo to mono.
torchaudio.info(path)
Sample rate, frames, channels — without decoding the file.
torchaudio.load(p, frame_offset=n, num_frames=k)
Read a slice. Essential once your clips are longer than memory.
AudioDecoder(p).get_samples_played_in_range(0, 5)
The native TorchCodec call, addressed in seconds rather than frames.

A3Spectrogram familywaveform → features

T.Spectrogram(n_fft=400, hop_length=None, power=2.0)★
STFT magnitude. power=2 energy, 1 amplitude, None keeps it complex.
T.MelSpectrogram(sample_rate=16000, n_mels=128)★
Literally Spectrogram then MelScale. The default input for most audio CNNs.
T.MFCC(sample_rate=16000, n_mfcc=40)★
DCT of the log-mel. Classic ASR features; pass melkwargs= to tune the mel stage.
T.AmplitudeToDB(stype="power")
Log scaling. Match stype to the power you used, or the dB range is wrong.
T.MelScale · T.InverseMelScale · T.LFCC
The stages on their own, plus the linear-filterbank cousin of MFCC.
T.GriffinLim() · T.InverseSpectrogram()
Back to audio. Griffin-Lim guesses the discarded phase.
spec.shape → (channel, freq, time)
freq = n_fft // 2 + 1, or n_mels for a mel spectrogram.
frames = 1 + samples // hop_length★
Output width. The one formula that tells you the time axis before you run anything.
n_fft sets resolution, hop_length sets overlap★
Bigger n_fft buys frequency detail and costs time detail. hop=n_fft//4 is the usual default.
T.Spectrogram(..., center=True, window_fn=torch.hann_window)
Padding and window. Change these and your features stop matching a pretrained model.

A4Resample & augmentSpecAugment

T.Resample(orig_freq=44100, new_freq=16000)★
Build it once and reuse — it precomputes a kernel, so constructing it per sample is slow.
T.FrequencyMasking(freq_mask_param=80)★
SpecAugment: blank out frequency bands. Operates on spectrograms, not waveforms.
T.TimeMasking(time_mask_param=80)
The time-axis counterpart. The pair is the standard ASR augmentation.
T.TimeStretch(rate, fixed_rate=True)
Needs a complex spectrogram — set power=None upstream.
T.PitchShift · T.Fade · T.Vol
Waveform-domain effects.
T.Vad() · T.SlidingWindowCmn() · T.ComputeDeltas()
Voice activity trimming, cepstral mean norm, delta features.
nn.Sequential(resample, spec, masking).to(dev)★
Transforms are Modules, so the whole feature pipeline runs on GPU and inside torch.compile.
resampler = T.Resample(...) once, reuse★
Build it outside the loop. The constructor precomputes a filter kernel; rebuilding per clip dominates your data pipeline.
T.AddNoise(waveform, noise, snr)
Additive noise at a target SNR — the most effective single augmentation for ASR robustness.
T.SpeedPerturbation(sr, [0.9, 1.0, 1.1])
The Kaldi-style three-way speed augmentation.

A5functionalstateless ops

import torchaudio.functional as F
Careful: the audio F shadows torch.nn.functional if you import both bare.
F.resample(w, orig, new)
One-shot version — fine for a single call, slow in a loop.
F.amplitude_to_DB · F.DB_to_amplitude
Scale conversions without constructing a Module.
F.lfilter(w, a_coeffs, b_coeffs)
Differentiable IIR filtering.
F.lowpass_biquad · highpass_biquad · bandpass_biquad
Ready-made biquads for pre-emphasis and band limiting.
F.forced_align(log_probs, targets)
CTC alignment — audio to transcript, frame by frame.
F.edit_distance(seq1, seq2)
The building block for WER and CER.
F.melscale_fbanks(...)
The mel filterbank matrix itself, if you are building features by hand.
F.preemphasis(w) · F.deemphasis(w)
Classic high-frequency tilt before feature extraction.
F.add_noise(w, n, snr) · F.loudness(w, sr)
Functional noise mixing; ITU-R BS.1770 loudness for normalization.

A6Pipelines & modelspretrained speech

bundle = torchaudio.pipelines.WAV2VEC2_ASR_BASE_960H★
Model + weights + labels + expected sample rate, in one object.
model = bundle.get_model()★
Downloads and instantiates, ready for .to(dev).
bundle.sample_rate · bundle.get_labels()★
Resample your audio to bundle.sample_rate first — usually 16 kHz. Feeding 44.1 kHz gives fluent nonsense.
emission, _ = model(waveform)
Frame-wise logits over the label set.
HUBERT_ASR_LARGE · MMS_FA · SQUIM_OBJECTIVE
Other bundles: ASR, forced alignment, speech quality.
models.decoder.ctc_decoder(...)
Beam search with an optional lexicon and language model.
models.Wav2Vec2Model · Conformer · Emformer
The bare architectures, when you want to train your own.
feats, _ = model.extract_features(waveform)★
Per-layer representations instead of ASR logits — this is how Wav2Vec2 is used as an encoder.
bundle.get_labels()
The token alphabet, index-aligned with the emission columns. Blank is index 0.

A7Datasets & batchingragged by nature

datasets.LIBRISPEECH(root, url="test-clean", download=True)★
The ASR default. Yields (waveform, sr, transcript, ...).
datasets.SPEECHCOMMANDS · YESNO · CMUARCTIC
Keyword spotting, a tiny smoke-test set, and TTS corpora.
collate_fn → nn.utils.rnn.pad_sequence★
Clips have different lengths, so the default collate will fail. Pad and return the true lengths.
sorted-by-length bucketing
Group similar durations per batch — far less wasted padding compute.
compliance.kaldi.fbank(w) · .mfcc(w)
Bit-compatible with Kaldi, for reproducing published recipes.
nn.utils.rnn.pad_sequence(seqs, batch_first=True)★
Ragged audio into a rectangle. Keep the true lengths — you need them next.
pack_padded_sequence(x, lens, batch_first=True, enforce_sorted=False)
Stops RNNs from consuming your padding. enforce_sorted=False saves you sorting the batch.

A8Speech to text, end to endthe whole path

bundle = pipelines.WAV2VEC2_ASR_BASE_960H★
Weights, expected sample rate and label set travel together as one object.
model = bundle.get_model().to(dev); labels = bundle.get_labels()
Downloads and caches on first call.
if sr != bundle.sample_rate: w = F.resample(w, sr, bundle.sample_rate)★silent
Never skip this. A wrong sample rate raises no error — it just returns confident nonsense.
with torch.inference_mode(): emission, _ = model(w)
Shape (batch, time, vocab): a log-probability per token per frame.
ids = emission[0].argmax(-1)
Greedy CTC decoding, step one — best token per frame.
ids = torch.unique_consecutive(ids)★
The CTC collapse. Merge repeats, then drop the blank (index 0). Miss this and you get “hheelllloo”.
"".join(labels[i] for i in ids if i != 0)
Indices back to characters. That is the transcript.
decoder = models.decoder.ctc_decoder(lexicon, tokens, lm=...)
Beam search with a language model when greedy is not accurate enough.

pytorch + torchvision · torchaudio

torch · the core

torchvision · images

torchaudio · sound

Four pictures worth the whole reference

1 · the autograd graph

2 · the shape rosetta

3 · who owns media I/O now

4 · the audio feature chain

Worth memorizing