PyTorch reference and troubleshooting

Use this page when you know the goal but need the next command or check. For the full reasoning, return to the core workflow or vision guide.

Short reminder

1. Shapes are part of the program: image batches normally use [N, C, H, W].
2. Model, inputs, targets, and helper tensors used together share a device.
3. A Dataset returns one sample; a DataLoader returns a batch.
4. A classifier normally returns raw logits.
5. The update order is clear → predict → measure → differentiate → update.
6. Evaluation uses both model.eval() and disabled gradients.
7. Convolution grows useful feature channels; pooling or stride reduces spatial size.
8. Training behavior matters only when compared with validation and real-use data.
9. Save weights together with the context needed to interpret them.

Choosing the next tool

The course's later image-model work adds choices to the same training loop. Use this as a prompt to investigate, not as a replacement for a controlled experiment.

When you need to…	First choice to consider	Check before trusting it
compare classifiers with uneven classes	precision, recall, F1, and a confusion matrix alongside accuracy	per-class results and the cost of each error type
change how fast a model learns	tune the optimizer's learning rate, then consider a scheduler	validation behavior at the same data split and run budget
search several settings	record trials and their validation objective, whether manual or automated	that the test set stays outside the search
reuse an image model	use the weights' recommended preprocessing and replace the classifier head	class order, input shape, and which parameters are trainable
choose between models	compare quality with latency, memory, and parameter count	measurements on the intended device

For connected explanations of these choices, open metrics and tuning, vision or text.

Course 2: what changes in practice

Course 2 routes and extra snippets

The same loop now supports more deliberate choices. Use these direct routes instead of re-reading a whole course:

Need to remember	Read
Precision, recall, macro F1, and metric state	Metrics
Step, cosine and plateau LR schedules	Schedulers
Search spaces, trials and TPE	Optuna
Salt and pepper noise or transform order	Augmentation
Weights metadata, boxes and masks	Pretrained vision
Frozen vs partially fine-tuned features	Vision stages
Tokenization, padding and vectors	Text representation
EmbeddingBag, class weights and DistilBERT	Text classifiers
Workers, prefetching and profiling	Efficient training
Mixed precision and accumulation	Precision · Accumulation

Metrics: accuracy can hide a weak minority class. Read per-class results and keep the final test set outside tuning.

Learning rate: the optimizer changes weights; a scheduler changes the optimizer's learning rate over time. For an epoch-based StepLR, call scheduler.step() after that epoch's training updates:

optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)
for epoch in range(epochs):
    train_one_epoch(model, train_loader, optimizer)  # your training function
    scheduler.step()  # after five calls, multiply the learning rate by 0.1

Transfer learning: reuse learned visual features and train a new final classifier first. The replacement layer must have one output per class; the optimizer should receive only parameters you intend to train:

from torchvision.models import resnet18, ResNet18_Weights
weights = ResNet18_Weights.DEFAULT
model = resnet18(weights=weights)
for parameter in model.parameters():
    parameter.requires_grad = False
model.fc = torch.nn.Linear(model.fc.in_features, num_classes)
optimizer = torch.optim.Adam(model.fc.parameters(), lr=1e-3)
preprocess = weights.transforms()  # use for the model's expected input format

The frozen backbone still participates in the forward pass; requires_grad=False prevents its weights from being updated. Decide on fine-tuning only after checking validation behavior. If you change which layers are trainable, rebuild the optimizer for the new parameter set.

Freezing parameters does not freeze BatchNorm statistics. Use a deliberate training-mode policy for the backbone.

Function finder

Each row groups related functions by the reason to use them. Follow the link for the shape contract, example and common mistake.

Tensors, models and the learning loop

Function / parameter	Remember	Explanation
tensor, from_numpy, zeros, ones, arange, rand	creation; copy vs shared storage	Tensor operations
reshape, flatten, squeeze, unsqueeze	change dimensions; preserve batch	Shapes
cat, stack, transpose, permute	join or reorder; not equivalent	Tensor operations
boolean masks, sum, mean, std, matmul	selection, reductions and matrix math	Tensor operations
.to, .float, .long, .item, .detach	device/dtype or Python scalar; gradient history	Dtype/device
Module, Sequential, ModuleList	register layers; define their flow	Flexible models
Linear, ReLU, Tanh, Sigmoid, Softmax	representation and output meaning	Activations and logits
MSELoss, L1Loss, SmoothL1Loss, CrossEntropyLoss, BCEWithLogitsLoss, NLLLoss	match task, target shape and output	Loss table
backward, zero_grad, step, SGD, Adam, AdamW	calculate then apply gradients	Update loop
train, eval, no_grad, inference_mode	module behavior vs gradient tracking	Evaluation
Dataset, Subset, random_split, DataLoader	samples, indices and batches	Data contracts
num_workers, pin_memory, prefetch_factor, persistent_workers, drop_last	CPU preparation and batching trade-offs	Loader settings
Conv2d, MaxPool2d, AdaptiveAvgPool2d	learned spatial features and reduction	CNN shapes
Dropout, BatchNorm2d, weight_decay	training-only variation / statistics / regularization	Hyperparameters
named_parameters, named_children, named_modules, register_forward_hook	inspect registered tensors and activations	Inspection

Images, pretrained models and metrics

Function / parameter	Remember	Explanation
decode_image, ToTensor, ToPILImage, make_grid, save_image	image type, range and visual inspection	Image utilities
ImageFolder, FakeData, EMNIST, SVHN	dataset-specific labels and split arguments	Dataset choices
Compose, Resize, CenterCrop, RandomResizedCrop, ColorJitter, RandomAffine	ordered, label-preserving transforms	Transform pipeline
custom __call__, Normalize	noise before normalization; fixed mean/std	Noise · Normalization
weights.transforms, weights.meta, topk, requires_grad_	pretrained input/label contract and adaptation	Pretrained models
draw_bounding_boxes, draw_segmentation_masks, argmax	display predictions; choose the correct axis	Task outputs
TorchMetrics update, compute, reset, average	aggregate across a measurement interval	Metrics
StepLR, CosineAnnealingLR, ReduceLROnPlateau	schedules and monitored direction	Schedulers
Optuna suggest_*, create_study, optimize, best_params, trials_dataframe	search → validation score → recorded trials	Optuna

Text and contextual models

Function / parameter	Remember	Explanation
AutoTokenizer, BertTokenizerFast, DataCollatorWithPadding	matching IDs, special tokens and padded batches	Tokenizer
Embedding, EmbeddingBag, collate_fn, offsets	vectors, pooling and variable-length inputs	Text classifiers
cosine_similarity, PCA	similarity and limited 2D visualization	Embedding interpretation
AutoModel, AutoModelForSequenceClassification, .logits, save_pretrained	contextual vectors vs class predictions	Text fine-tuning

Training tools and efficiency

Function / parameter	Remember	Explanation
LightningModule, LightningDataModule, Trainer, self.log	separate learning logic and orchestration	Lightning
EarlyStopping, ModelCheckpoint, fast_dev_run, callback hooks	stopping, best weights and diagnostics	Callbacks
PyTorchProfiler, schedule, profile_memory, record_shapes	identify expensive work	Profiling
autocast, GradScaler, precision, accumulate_grad_batches	numerical precision and update boundaries	Efficient training
synchronize, reset_peak_memory_stats, max_memory_allocated, element_size	measure completion and tensor bytes	Measurements
state_dict, save, load, load_state_dict	retain weights plus their interpretation	Saving

Cheatsheet

Inspect data and models

print(x.shape, x.dtype, x.device)
print(x.min().item(), x.max().item())
print(labels.unique())
print(model)

Device

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
inputs, targets = inputs.to(device), targets.to(device)

Train

model.train()
optimizer.zero_grad(set_to_none=True)
logits = model(inputs)
loss = loss_fn(logits, targets)
loss.backward()
optimizer.step()

Evaluate

model.eval()
with torch.no_grad():
    logits = model(inputs)
    predictions = logits.argmax(dim=1)

Save and restore

torch.save(model.state_dict(), "model.pth")
state = torch.load("model.pth", map_location=device, weights_only=True)
model.load_state_dict(state)
model.eval()

Count trainable parameters

trainable = sum(
    parameter.numel()
    for parameter in model.parameters()
    if parameter.requires_grad
)

Fix common random seeds

import random
import numpy as np
import torch

random.seed(42)
np.random.seed(42)
torch.manual_seed(42)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(42)

Common errors

mat1 and mat2 shapes cannot be multiplied

The final input dimension does not match nn.Linear(in_features, ...).

print(x.shape)  # immediately before the linear layer

Preserve the batch when flattening: torch.flatten(x, start_dim=1).

Expected all tensors to be on the same device

Move the model, inputs, targets, and newly created helper tensors to the same device. Inspect .device at the failing operation.

Target N is out of bounds

For K output classes, labels for CrossEntropyLoss must normally be integer ids in 0..K-1. Check one-based source labels and the class mapping.

Loss does not improve

1. Inspect input values, labels, shapes, and dtypes.
2. Verify the output/loss pairing.
3. Try to overfit one small batch.
4. Confirm parameters receive gradients.
5. Confirm the optimizer owns those parameters.
6. Inspect the learning rate.

Validation changes unexpectedly

• Call model.eval() and disable gradients.
• Remove random validation transforms.
• Use a fixed validation split.
• Keep class mapping and deterministic preprocessing consistent.
• Check for overlap or entity leakage between splits.

CUDA out of memory

• Reduce batch size first.
• Do not retain computation graphs in Python lists.
• Store loss.item() rather than the loss tensor.
• Evaluate inside torch.no_grad().
• Reduce input resolution or model size after measuring the bottleneck.

NaN loss

• Inspect inputs for NaN or infinity.
• Reduce the learning rate.
• Check logarithms, divisions, and normalization.
• Confirm target dtypes and ranges.
• Clip gradients only when they genuinely explode.

Debug in this order

1. Print shapes, dtypes, devices, and value ranges.
2. Check labels and class range.
3. Verify model output shape and loss pairing.
4. Confirm train/eval mode.
5. Overfit one tiny batch.
6. Inspect learning curves and class-specific errors.
7. Only then change architecture or regularization.

Glossary

Activation

The output produced by a layer or nonlinear function.

Autograd

PyTorch's automatic differentiation system.

Batch

A group of samples processed together. With accumulation, several microbatches contribute to one optimizer update.

Epoch

One pass through the training dataset.

Microbatch

A smaller batch processed in one forward/backward pass. Several can accumulate gradients before one update.

Optimizer step

One call that updates the optimizer's owned parameters using their current gradients. It need not happen after every microbatch.

Feature map

One channel of activations produced by a convolutional filter.

Gradient

The derivative of the objective with respect to a parameter.

Logit

A raw model score before conversion to a probability.

Loss

A differentiable scalar objective measuring prediction error.

Parameter

A registered model tensor such as a weight or bias. It remains a parameter when frozen; requires_grad determines whether gradients are recorded for it.

Regularization

A technique intended to improve generalization rather than only training fit.

Tensor

A multidimensional array carrying shape, dtype, and device information.