PyTorch Training Loop

This note gathers the full, complete PyTorch training loop — every concept from this category so far (autograd, nn.Module, DataLoader, optimizers) assembled into the exact runnable pattern used throughout this entire hub, with a focus on the practical PyTorch-specific gotchas.

The Complete Loop

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MyModel().to(device)
optimizer = optim.AdamW(model.parameters(), lr=1e-4)
loss_fn = nn.CrossEntropyLoss()
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0

    for x_batch, y_batch in train_loader:
        x_batch, y_batch = x_batch.to(device), y_batch.to(device)   # move data to the SAME device as the model

        optimizer.zero_grad()
        predictions = model(x_batch)
        loss = loss_fn(predictions, y_batch)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()   # .item() extracts a Python float -- avoids accumulating tensors

    print(f"Epoch {epoch+1}: loss={running_loss/len(train_loader):.4f}")

Why .item() Matters for Logging

Accumulating loss (a tensor still attached to the computational graph) directly, rather than loss.item(), would keep the entire graph alive for every batch across the whole epoch — a serious, easily-overlooked memory leak. .item() extracts just the scalar Python float value, safely detached from any graph.

The Device-Mismatch Error

# A classic runtime error:
# RuntimeError: Expected all tensors to be on the same device, but found at least two devices

# The fix: EVERY tensor involved in a computation -- model, input, labels -- must be on the SAME device
model = model.to(device)
x_batch = x_batch.to(device)
y_batch = y_batch.to(device)

Common Mistakes

• Accumulating loss directly instead of loss.item() for logging — causes graph memory to accumulate across an entire epoch unnecessarily.
• Moving the model to GPU but forgetting to move the data (or vice versa) — produces the classic "tensors on different devices" runtime error.
• Forgetting model.train() at the start of each epoch, especially right after a validation pass that called model.eval().

Interview Relevance

Q: "Why should you call .item() when accumulating loss values for logging, rather than summing the raw loss tensor directly?" The raw loss tensor remains attached to its computational graph — accumulating it directly across many batches keeps every one of those graphs alive in memory unnecessarily, since nothing ever triggers their release. .item() extracts just the scalar value as a plain Python float, fully detached from the graph, avoiding this memory leak.

Practice Question

What's the minimum set of tensors that must all be moved to the same device (e.g. GPU) for a training step to run without a device-mismatch error?

Related DL Notes

• Training Loop
• PyTorch GPU/CUDA
• PyTorch Validation Loop