PyTorch 2 cheatsheet

PyTorch 2: Training and saving

6 common PyTorch 2 tasks.

Write a training loop

The standard loop: forward pass, loss, backward pass, update.

model.train()
for x, y in loader:
    x, y = x.to(device), y.to(device)
    optimizer.zero_grad()
    loss = loss_fn(model(x), y)
    loss.backward()
    optimizer.step()

Choose a loss and optimizer

Set up cross-entropy loss and the Adam optimizer.

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

Evaluate a model

Switch to evaluation mode and turn off gradients.

model.eval()
correct = 0
with torch.no_grad():
    for x, y in test_loader:
        preds = model(x.to(device)).argmax(dim=1)
        correct += (preds == y.to(device)).sum().item()

Save a model

Save the weights, which is recommended over pickling the whole model.

torch.save(model.state_dict(), "model.pt")

Load a model

Create the model, then load the saved weights.

model = Net()
model.load_state_dict(torch.load("model.pt", weights_only=True))
model.eval()

Schedule the learning rate

Reduce the learning rate over time. Call scheduler.step() once per epoch.

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)