Exercise: Debug a Vision Training Pipeline

Build the mental model

Bugs in a vision pipeline come in two flavors: the kind that crashes immediately with a shape error, and the kind that runs to completion while quietly producing a worse model. A shape mismatch — where a Flatten operation doesn't match what the first Linear layer expects — is the loud kind, and the traceback itself tells you where to look. But mode mix-ups between model.train() and model.eval(), applying random augmentation to validation data, or normalizing with numbers that don't match the stated preprocessing scheme are the quiet kind: nothing throws, the loss curve looks plausible, and you only notice something is wrong because accuracy plateaus lower than it should or validation numbers don't track training numbers the way they normally do.

Connect it to a real scenario

The Tutorial Platform's 'Code Review Mode' lets you annotate this exact script line by line — flag each suspected bug with a comment and attach a short 'why this breaks' explanation, then compare your annotations side by side against the platform's reference solution notes.

Shape-mismatch bugs throw a runtime error immediately, so the traceback tells you exactly where to look.
Mode mix-ups like train()/eval() confusion and augmentation applied to validation data corrupt results silently, with no error at all.
Always cross-check normalization statistics against whatever preprocessing scheme the code claims to use.

Try the working example

python

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision.transforms import v2
from torch.utils.data import Dataset, DataLoader

torch.manual_seed(0)

# --- Preprocessing: "ImageNet normalization" ---
train_transform = v2.Compose([
    v2.RandomHorizontalFlip(),
    v2.RandomCrop(224, padding=4),
    v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
])

val_transform = v2.Compose([
    v2.RandomHorizontalFlip(),
    v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
])

class SyntheticImageDataset(Dataset):
    # Stands in for a real folder of labeled photos (e.g. torchvision's
    # ImageFolder) so this script is fully self-contained and runnable.
    def __init__(self, num_samples, num_classes, transform):
        self.num_samples = num_samples
        self.num_classes = num_classes
        self.transform = transform

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        image = torch.rand(3, 224, 224)  # fake RGB photo
        label = torch.randint(0, self.num_classes, (1,)).item()
        return self.transform(image), label

train_set = SyntheticImageDataset(num_samples=64, num_classes=10, transform=train_transform)
val_set = SyntheticImageDataset(num_samples=32, num_classes=10, transform=val_transform)

train_loader = DataLoader(train_set, batch_size=32, shuffle=True)
val_loader = DataLoader(val_set, batch_size=32, shuffle=False)

class SimpleCNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        # 224x224 input -> two 2x2 pools -> 56x56 feature map, 32 channels
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(32 * 28 * 28, 128),
            nn.ReLU(),
            nn.Linear(128, num_classes),
        )

    def forward(self, x):
        x = self.features(x)
        x = self.classifier(x)
        return x

model = SimpleCNN(num_classes=train_set.num_classes)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

def train_one_epoch():
    for images, labels in train_loader:
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

def evaluate():
    model.eval()
    correct, total = 0, 0
    with torch.no_grad():
        for images, labels in val_loader:
            outputs = model(images)
            preds = outputs.argmax(dim=1)
            correct += (preds == labels).sum().item()
            total += labels.size(0)
    return correct / total

for epoch in range(5):
    train_one_epoch()
    acc = evaluate()
    print(f"Epoch {epoch+1}: val accuracy = {acc:.4f}")

You should see

Running this script fails immediately with a shape error: the Linear layer is sized for a 28x28 feature map (`32 * 28 * 28`), but two 2x2 max-pools on a 224x224 input actually produce 56x56, so PyTorch raises `RuntimeError: mat1 and mat2 shapes cannot be multiplied`. Fix that and a second bug surfaces silently — `optimizer.zero_grad()` is never called, so gradients accumulate across batches instead of resetting each step, and the model trains poorly (loss decreases erratically or plateaus high, with no error at all). The third bug, RandomHorizontalFlip applied inside val_transform, doesn't crash anything either — it just makes validation accuracy noisier and slightly non-reproducible run to run, since the model is being evaluated on randomly flipped images instead of a fixed validation set.

5-minute try-it

Before running the code, read it closely and identify all three planted bugs. For each one, explain where it is, why it's wrong, and whether it would throw an immediate error or silently corrupt training.

One important caution

Fixing only the shape error and declaring victory, missing the two bugs that don't produce any error message at all.

Dismissing the missing zero_grad() call as a minor style issue rather than a real bug that corrupts gradient accumulation and destroys training dynamics.

PyTorch Tutorials — Training a Classifier (CIFAR10) — Computer Vision