pytorch-guide

PyTorch Guide

Overview

PyTorch is the dominant deep learning framework in academic research, used in the majority of papers at NeurIPS, ICML, and ICLR. Its eager execution model, Pythonic API, and seamless integration with the Python scientific stack make it the default choice for prototyping and publishing research code.

However, PyTorch's flexibility is a double-edged sword. Subtle bugs -- forgetting model.eval(), accumulating gradients across batches, incorrect device placement, memory leaks from detached tensors -- can silently corrupt results without raising errors. These issues are especially dangerous in research settings where ground truth is unknown.

This guide catalogs the most common PyTorch mistakes, provides battle-tested training patterns, and covers performance optimization techniques that every researcher should know. The patterns here are drawn from top-tier ML research codebases and the PyTorch team's own best practice recommendations.

Common Mistakes and Fixes

The Big Five Mistakes

# MISTAKE 1: Forgetting model.eval() and torch.no_grad()
# This causes dropout and batch norm to behave incorrectly during evaluation
# and wastes memory by tracking gradients

# WRONG
def evaluate(model, dataloader):
    total_correct = 0
    for x, y in dataloader:
        output = model(x)  # Dropout still active! BN using batch stats!
        total_correct += (output.argmax(1) == y).sum().item()

# RIGHT
@torch.no_grad()
def evaluate(model, dataloader):
    model.eval()
    total_correct = 0
    for x, y in dataloader:
        output = model(x)
        total_correct += (output.argmax(1) == y).sum().item()
    model.train()  # Restore training mode
    return total_correct

# MISTAKE 2: Not zeroing gradients (they accumulate by default!)
# WRONG - gradients from previous batch add to current batch
for x, y in dataloader:
    loss = criterion(model(x), y)
    loss.backward()
    optimizer.step()

# RIGHT
for x, y in dataloader:
    optimizer.zero_grad()        # Clear previous gradients
    loss = criterion(model(x), y)
    loss.backward()
    optimizer.step()

# BETTER (slightly faster, avoids memset)
for x, y in dataloader:
    optimizer.zero_grad(set_to_none=True)
    loss = criterion(model(x), y)
    loss.backward()
    optimizer.step()