The PyTorch training loop

1. Forward pass — compute predictions by calling model(X).
2. Loss — compare predictions vs targets with a loss function.
3. Backward pass — call loss.backward() to compute gradients via autodiff.
4. Optimizer step — call optimizer.step() to update parameters.
5. Zero gradients — call optimizer.zero_grad() before the next batch.

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MyModel().to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-2)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)

for epoch in range(num_epochs):
    model.train()
    for X, y in train_loader:
        X, y = X.to(device), y.to(device)
        optimizer.zero_grad()
        pred = model(X)
        loss = loss_fn(pred, y)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)  # gradient clipping
        optimizer.step()
    scheduler.step()

    model.eval()
    with torch.no_grad():
        val_loss = sum(loss_fn(model(X.to(device)), y.to(device)) for X, y in val_loader)
    print(f"Epoch {epoch}: val_loss={val_loss:.4f}")

Gradient clipping (clip_grad_norm_) prevents exploding gradients — especially important when training RNNs or large transformers. The cosine LR scheduler gradually cools the learning rate over training.

Key abstractions

x = torch.randn(3, 4, device='cuda')
x.half()   # float16

with torch.no_grad():   # disable grad tracking
    pred = model(X)

class Net(nn.Module):
    def forward(self, x): ...

DataLoader(ds, batch_size=64,
  num_workers=4, pin_memory=True)

optim.AdamW(model.parameters(),
  lr=3e-4, weight_decay=1e-2)

loss = nn.CrossEntropyLoss()
loss(logits, targets)

Building custom nn.Modules

Every model is a subclass of nn.Module. Define parameters in __init__, forward logic in forward. Compose modules freely — PyTorch handles parameter registration and device movement automatically.

class ResBlock(nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, dim),
            nn.LayerNorm(dim),
            nn.GELU(),
            nn.Linear(dim, dim),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x + self.net(x)   # residual connection

class MLP(nn.Module):
    def __init__(self, in_dim: int, out_dim: int, depth: int = 4):
        super().__init__()
        self.stem = nn.Linear(in_dim, 256)
        self.blocks = nn.ModuleList([ResBlock(256) for _ in range(depth)])
        self.head = nn.Linear(256, out_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.stem(x)
        for block in self.blocks:
            x = block(x)
        return self.head(x)

Data pipeline: Dataset + DataLoader

Implement two methods on Dataset: __len__ and __getitem__. PyTorch handles the rest. For image data, use torchvision.transforms inside __getitem__ for on-the-fly augmentation.

from torch.utils.data import Dataset, DataLoader

class TabularDataset(Dataset):
    def __init__(self, X: np.ndarray, y: np.ndarray):
        self.X = torch.from_numpy(X).float()
        self.y = torch.from_numpy(y).float()

    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]

train_loader = DataLoader(
    TabularDataset(X_train, y_train),
    batch_size=256,
    shuffle=True,
    num_workers=4,        # parallel CPU workers
    pin_memory=True,      # faster CPU→GPU transfer
)

pin_memory=True keeps tensors in pinned (page-locked) CPU memory, enabling faster async transfers to the GPU. Always profile with torch.profiler if training is slower than expected — DataLoader workers are the most common bottleneck.

Mixed precision & performance tips

scaler = torch.cuda.amp.GradScaler()
with torch.autocast(device_type='cuda'):
    loss = loss_fn(model(X), y)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

model = torch.compile(model)
# rest of training loop unchanged

model.eval()
with torch.no_grad():
    preds = model(X_test)

from torch.utils.checkpoint import checkpoint
out = checkpoint(layer, x)  # recompute on backward

Saving, loading & exporting

Save checkpoints with both model and optimizer state so training can resume. For inference serving, export to ONNX or use torch.export (PyTorch 2.x) for a portable, framework-agnostic representation.

# Save
torch.save({
    "epoch": epoch,
    "model_state_dict": model.state_dict(),
    "optimizer_state_dict": optimizer.state_dict(),
    "loss": best_loss,
}, "checkpoint.pt")

# Load
checkpoint = torch.load("checkpoint.pt", map_location=device)
model.load_state_dict(checkpoint["model_state_dict"])
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])

# Export to ONNX for serving
dummy_input = torch.randn(1, in_dim, device=device)
torch.onnx.export(model, dummy_input, "model.onnx", opset_version=17)

Debugging & inspecting models