A checkpoint is a snapshot of model.state_dict() (all learnable parameters + buffers like BatchNorm running stats) plus optimizer.state_dict() (momentum/variance estimates for Adam, etc.) and any training metadata (epoch, best validation score, RNG state). Mixed precision (AMP - Automatic Mixed Precision) runs most ops in float16/bfloat16 while keeping a master copy of weights and precision-sensitive ops (reductions, loss computation) in float32, using a GradScaler to prevent gradient underflow in float16.

# ---- Saving a checkpoint ----
checkpoint = {
    "epoch": epoch,
    "model_state_dict": model.state_dict(),
    "optimizer_state_dict": optimizer.state_dict(),
    "best_val_loss": best_val_loss,
}
torch.save(checkpoint, f"checkpoints/model_epoch{epoch}.pt")

# ---- Loading and resuming ----
checkpoint = torch.load("checkpoints/model_epoch10.pt", map_location=device)
model.load_state_dict(checkpoint["model_state_dict"])
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
start_epoch = checkpoint["epoch"] + 1
best_val_loss = checkpoint["best_val_loss"]

model.to(device)
for epoch in range(start_epoch, num_epochs):
    ...  # continue training exactly where it left off

torch.save(model.state_dict(), ...) is preferred over torch.save(model, ...) (pickling the whole object) because it's robust to refactors in the model's class definition, doesn't tie the checkpoint to exact file paths/class locations, and is the portable format expected by most downstream tooling. Always pass map_location=device when loading, so a checkpoint saved on GPU can still be loaded on a CPU-only machine.

Loading safely: since PyTorch 2.6, torch.load defaults to weights_only=True, which refuses to unpickle arbitrary Python objects (pickles can execute code). Keep checkpoints to tensors and plain Python types - as above - and load untrusted weights only from safetensors files (see PyTorch for LLMs).

Saving only the best model (common pattern to avoid disk bloat from saving every epoch):

if val_loss < best_val_loss:
    best_val_loss = val_loss
    torch.save(model.state_dict(), "checkpoints/best_model.pt")

from torch.amp import autocast, GradScaler   # torch.cuda.amp is deprecated

model = MyModel().to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
scaler = GradScaler("cuda")                  # only needed for float16, not bfloat16

for epoch in range(num_epochs):
    model.train()
    for batch_x, batch_y in train_loader:
        batch_x, batch_y = batch_x.to(device), batch_y.to(device)
        optimizer.zero_grad()

        with autocast(device_type="cuda", dtype=torch.float16):  # ops inside run in fp16 where safe
            outputs = model(batch_x)
            loss = loss_fn(outputs, batch_y)

        scaler.scale(loss).backward()          # scale loss up before backward - avoids fp16 underflow
        scaler.step(optimizer)                  # unscales gradients, then calls optimizer.step()
        scaler.update()                         # adjusts the scale factor for next iteration

Why this saves VRAM and speeds up training:

• Memory: activations and gradients stored in float16 (2 bytes) instead of float32 (4 bytes) roughly halves activation memory - the single biggest VRAM consumer during training for large batch sizes/sequence lengths
• Speed: modern GPU tensor cores (NVIDIA Volta and newer) execute float16/bfloat16 matrix multiplications at significantly higher throughput than float32
• autocast() automatically chooses which ops are safe to run in reduced precision (matmuls, convolutions) versus which need float32 for numerical stability (softmax, loss reductions, batch norm statistics) - you don't have to manually cast tensors
• bfloat16 (available on newer GPUs/TPUs) has the same exponent range as float32, so it often doesn't need GradScaler at all - it trades some mantissa precision for range stability, avoiding the underflow problem float16 has