PyTorch Cheat Sheet
Essential PyTorch syntax for tensors, autograd, building neural network modules, and writing a standard training loop for deep learning models.
2 PagesIntermediateApr 10, 2026
Tensor Basics
Creating and manipulating tensors.
python
import torchx = torch.tensor([[1.0, 2.0], [3.0, 4.0]])y = torch.zeros(2, 2)z = torch.rand(2, 2)device = "cuda" if torch.cuda.is_available() else "cpu"x = x.to(device)a = x + y # elementwise addb = x @ y # matrix multiplyc = x.view(-1, 4) # reshape (view shares memory)print(x.shape, x.dtype)
Autograd
Automatic differentiation for gradients.
python
x = torch.tensor(2.0, requires_grad=True)y = x ** 2 + 3 * xy.backward() # compute dy/dxprint(x.grad) # tensor(7.) since dy/dx = 2x + 3with torch.no_grad(): # disable grad tracking (inference) z = x * 2
Model & Training Loop
Define a network and train it.
python
import torch.nn as nnimport torch.optim as optimclass Net(nn.Module): def __init__(self): super().__init__() self.fc1 = nn.Linear(784, 128) self.fc2 = nn.Linear(128, 10) def forward(self, x): x = torch.relu(self.fc1(x)) return self.fc2(x)model = Net().to(device)optimizer = optim.Adam(model.parameters(), lr=1e-3)criterion = nn.CrossEntropyLoss()for epoch in range(epochs): for inputs, labels in dataloader: optimizer.zero_grad() # clear gradients outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() # backprop optimizer.step() # update weights
Common Layers & Losses
Frequently used nn.Module building blocks.
- nn.Linear- fully connected layer
- nn.Conv2d- 2D convolution for image data
- nn.LSTM- recurrent layer for sequence data
- nn.Dropout- regularization by zeroing random activations
- nn.BatchNorm2d- normalizes activations across the batch
- nn.CrossEntropyLoss- combines LogSoftmax + NLLLoss for classification
- nn.MSELoss- mean squared error for regression
- torch.optim.Adam / SGD- optimizers that update parameters from gradients
Custom Autograd Function
Define a custom forward/backward pass by subclassing torch.autograd.Function.
python
class ClampGrad(torch.autograd.Function): @staticmethod def forward(ctx, x, min_val, max_val): ctx.save_for_backward(x) ctx.min_val, ctx.max_val = min_val, max_val return x.clamp(min_val, max_val) @staticmethod def backward(ctx, grad_output): (x,) = ctx.saved_tensors mask = (x >= ctx.min_val) & (x <= ctx.max_val) return grad_output * mask, None, Noney = ClampGrad.apply(x, -1.0, 1.0)y.sum().backward()
Custom Dataset & DataLoader Tuning
Implement a Dataset and tune the DataLoader for throughput.
python
from torch.utils.data import Dataset, DataLoaderclass TensorDataset(Dataset): def __init__(self, X, y): self.X, self.y = X, y def __len__(self): return len(self.X) def __getitem__(self, idx): return self.X[idx], self.y[idx]loader = DataLoader( TensorDataset(X, y), batch_size=64, shuffle=True, num_workers=4, # parallel worker processes pin_memory=True, # faster host->GPU transfer persistent_workers=True, drop_last=True,)
Mixed Precision Training (AMP)
Speed up training and reduce memory with automatic mixed precision.
python
scaler = torch.cuda.amp.GradScaler()for inputs, labels in dataloader: inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad(set_to_none=True) with torch.autocast(device_type="cuda", dtype=torch.float16): outputs = model(inputs) loss = criterion(outputs, labels) scaler.scale(loss).backward() # scales loss to avoid underflow scaler.step(optimizer) scaler.update()
Checkpointing & DistributedDataParallel
Save/resume full training state and scale training across GPUs.
python
# Save full training state, not just weightstorch.save({ "epoch": epoch, "model_state": model.state_dict(), "optimizer_state": optimizer.state_dict(), "scheduler_state": scheduler.state_dict(),}, "checkpoint.pt")ckpt = torch.load("checkpoint.pt", map_location=device)model.load_state_dict(ckpt["model_state"])optimizer.load_state_dict(ckpt["optimizer_state"])# Multi-GPU training entry pointimport torch.distributed as distfrom torch.nn.parallel import DistributedDataParallel as DDPdist.init_process_group(backend="nccl")model = DDP(model.to(local_rank), device_ids=[local_rank])
Advanced APIs & Gotchas
Lesser-known utilities and common pitfalls beyond the basics.
- torch.compile(model)- JIT-compiles a model graph for significant speedups on PyTorch 2.x
- register_hook / register_forward_hook- intercept gradients or activations for debugging without changing forward code
- torch.utils.checkpoint- trades compute for memory by recomputing activations during backward
- nn.utils.parametrize- reparameterize weights (e.g. weight normalization) without subclassing layers
- in-place ops on leaf tensors- e.g. x += 1 on a tensor with requires_grad=True raises a RuntimeError during backward
- tensor.detach()- returns a new tensor sharing storage but detached from the autograd graph
- torch.einsum- concise Einstein-summation notation for complex tensor contractions
- non_blocking=True- combined with pin_memory to overlap host-to-device copies with compute
Pro Tip
Call model.eval() and wrap inference in torch.no_grad() to disable dropout/batchnorm training behavior and gradient tracking — forgetting this is a common source of inconsistent validation metrics.
Was this cheat sheet helpful?
Explore Topics
#PyTorch#PyTorchCheatSheet#DataScience#Intermediate#TensorBasics#Autograd#ModelTrainingLoop#CommonLayersLosses#Networking#MachineLearning#CheatSheet#SkillVeris
