PyTorch Tips and Best Practices

These PyTorch tips and best practices collect the habits that save the most time in real projects. Most training bugs are silent: the code runs, but the model learns poorly. A few disciplined checks catch them early.

Make Runs Reproducible

import random, numpy as np, torch

def set_seed(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)   # also seeds all CUDA devices

set_seed(42)
print(torch.rand(2))

tensor([0.8823, 0.9150])

For fully deterministic GPU results you can also call torch.use_deterministic_algorithms(True), at some cost in speed.

Overfit One Batch First

Before a long training run, check that your model can drive the loss near zero on a single small batch. If it cannot, there is a bug in the model, loss or data, not a tuning problem.

from torch import nn

model = nn.Sequential(nn.Linear(20, 64), nn.ReLU(), nn.Linear(64, 3))
xb, yb = torch.randn(16, 20), torch.randint(0, 3, (16,))
opt = torch.optim.AdamW(model.parameters(), lr=1e-2)
for step in range(200):
    loss = nn.functional.cross_entropy(model(xb), yb)
    opt.zero_grad(); loss.backward(); opt.step()
print(round(loss.item(), 4))

0.0012   (sample output — should be close to 0)

Debug Shapes Early

Print shapes at each step of forward while developing, or pass a dummy batch through the model:

x = torch.randn(4, 20)
for layer in model:
    x = layer(x)
    print(layer.__class__.__name__, tuple(x.shape))

Linear (4, 64)
ReLU (4, 64)
Linear (4, 3)

Gradient Clipping

Clipping the gradient norm prevents occasional huge updates that make the loss explode, and is standard when training transformers and RNNs:

loss = nn.functional.cross_entropy(model(xb), yb)
opt.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
opt.step()

(gradients rescaled so their total norm is at most 1.0)

Speed Up with torch.compile

compiled = torch.compile(model)
out = compiled(torch.randn(8, 20))
print(out.shape)

torch.Size([8, 3])

The first call is slow while PyTorch compiles optimized kernels; later calls are faster. Gains are largest on GPUs with bigger models.

Common Mistakes to Avoid

• Forgetting optimizer.zero_grad(), so gradients accumulate across steps.
• Forgetting model.eval() and torch.no_grad() during validation.
• Adding softmax before CrossEntropyLoss.
• Labels with the wrong dtype: CrossEntropyLoss needs int64 class indices.
• Accumulating loss tensors in a list (keeps the graph alive); use loss.item().
• Not shuffling training data, or shuffling validation data and comparing outputs to labels in the wrong order.
• Mixing devices: some tensors on CPU, others on GPU.

Where to Go Next

You now know the core of PyTorch. Good next steps are learning about transformers and LLMs with Hugging Face, trying PyTorch Lightning for less boilerplate, logging experiments with TensorBoard or Weights and Biases, and reading the source of well-known models in torchvision.

Frequently asked questions

How do I make PyTorch results reproducible?

Seed Python, NumPy and PyTorch with the same value, avoid nondeterministic operations or enable torch.use_deterministic_algorithms(True), and fix DataLoader shuffling with a seeded generator.

Is torch.compile worth using?

Often yes for training and inference on GPUs, where it can give significant speedups with one line of code. Measure on your own model, as gains vary.

Why is my PyTorch model not learning?

Check that it can overfit one batch, that the learning rate is reasonable, that inputs are normalized, that labels match the loss function, and that zero_grad and step are called.