How do you save and load PyTorch models correctly, including full training checkpoints?

import torch
import torch.nn as nn

model     = nn.Linear(10, 5)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# ── RECOMMENDED: save/load state_dict only
torch.save(model.state_dict(), "weights.pt")

model_new = nn.Linear(10, 5)        # must define the SAME architecture first
model_new.load_state_dict(torch.load("weights.pt"))
model_new.eval()                     # always call before inference

# ── NOT recommended: save the entire model object
# Fragile — breaks if the class definition moves or changes
torch.save(model, "full_model.pt")
loaded_model = torch.load("full_model.pt", weights_only=False)

# ── Full training checkpoint — for resuming training
def save_checkpoint(path, epoch, model, optimizer, best_val_loss):
    torch.save({
        "epoch":            epoch,
        "model_state":      model.state_dict(),
        "optimizer_state":  optimizer.state_dict(),  # Adam momentum buffers etc.
        "best_val_loss":    best_val_loss,
    }, path)

def load_checkpoint(path, model, optimizer):
    ckpt = torch.load(path, map_location="cpu")  # always load to CPU first
    model.load_state_dict(ckpt["model_state"])
    optimizer.load_state_dict(ckpt["optimizer_state"])
    return ckpt["epoch"], ckpt["best_val_loss"]

save_checkpoint("ckpt.pt", epoch=5, model=model, optimizer=optimizer, best_val_loss=0.42)
epoch, best_loss = load_checkpoint("ckpt.pt", model_new, optimizer)

# ── Loading on a different device than it was saved
model.load_state_dict(
    torch.load("weights.pt", map_location="cpu")  # avoid GPU OOM if GPU unavailable
)
model = model.to("cuda")  # then move to the desired device