Every PyTorch project has the same heartbeat: for each batch — forward, loss, zero_grad, backward, step. Around it: an evaluation loop in eval() mode under no_grad(), and checkpoints.

model.train()
Dropout on, BatchNorm uses batch statistics.
Forward
logits = model(xb) then loss = loss_fn(logits, yb).
optimizer.zero_grad()
Clear gradients from the previous step.
loss.backward()
Compute gradients.
optimizer.step()
Update weights (then scheduler.step() if you use one).
Evaluate
model.eval() + torch.no_grad() on the validation set; keep the best checkpoint.

A complete example

Classify points by whether they lie inside a circle — a non-linear boundary a linear model can't learn.

Python

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

torch.manual_seed(0)
X = torch.rand(1000, 2) * 2 - 1                      # points in [-1, 1]²
y = ((X ** 2).sum(dim=1) < 0.5).long()               # class 1 inside the circle
train_ds, val_ds = TensorDataset(X[:800], y[:800]), TensorDataset(X[800:], y[800:])
train_loader = DataLoader(train_ds, batch_size=32, shuffle=True)
val_loader = DataLoader(val_ds, batch_size=200)

model = nn.Sequential(nn.Linear(2, 32), nn.ReLU(), nn.Linear(32, 32), nn.ReLU(), nn.Linear(32, 2))
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-2)

def evaluate(loader):
    model.eval()
    correct = total = 0
    with torch.no_grad():
        for xb, yb in loader:
            correct += (model(xb).argmax(dim=1) == yb).sum().item()
            total += len(yb)
    return correct / total

for epoch in range(20):
    model.train()
    for xb, yb in train_loader:
        loss = loss_fn(model(xb), yb)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

print(f"train acc {evaluate(train_loader):.2f}, val acc {evaluate(val_loader):.2f}")   # train acc 0.99, val acc 0.97

Saving and loading

Python

import torch
from torch import nn

torch.manual_seed(0)
model = nn.Linear(3, 1)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

torch.save(model.state_dict(), "model.pt")                   # weights only — the usual way
torch.save({"epoch": 5, "model": model.state_dict(),          # full checkpoint to resume training
            "optimizer": optimizer.state_dict()}, "checkpoint.pt")

restored = nn.Linear(3, 1)                                    # rebuild the architecture first…
restored.load_state_dict(torch.load("model.pt"))              # …then load the weights
restored.eval()
x = torch.ones(1, 3)
print(torch.equal(model(x), restored(x)))                     # True

ckpt = torch.load("checkpoint.pt")
print(ckpt["epoch"], list(ckpt["model"].keys()))              # 5 ['weight', 'bias']

Using a GPU

Python

device = torch.device("cuda" if torch.cuda.is_available()
                      else "mps" if torch.backends.mps.is_available() else "cpu")
model = model.to(device)                     # parameters move once
for xb, yb in train_loader:
    xb, yb = xb.to(device), yb.to(device)    # every batch must follow
    ...
# load a GPU checkpoint on a CPU-only machine:
state = torch.load("model.pt", map_location="cpu")

Walk me through a PyTorch training step.

What does model.eval() change?