PyTorch Saving, Loading, Exporting and Deployment

Checkpoint Contents

Training is not the finish line. A model must be saved, loaded, versioned, and served consistently. Inference code should use the same preprocessing as training, put the model in eval mode, disable gradients, and return stable outputs.

PyTorch commonly saves state dictionaries rather than entire model objects. This is safer because code defines architecture and checkpoint files store learned weights.

A training checkpoint should contain enough information to resume training and audit the experiment. An inference artifact may contain only the model state and class labels.

Save model_state_dict and optimizer_state_dict for resumable training.
Save class names and preprocessing configuration for inference.
Store model version, git commit, data version, and metrics in metadata.

Inference Safety

Inference should be deterministic and memory efficient. Use eval mode, no_grad or inference_mode, correct device handling, and input validation.

Never train accidentally in an API handler.
Validate input shape and dtype before prediction.
Return probabilities or class labels based on product needs.

Save and Load a Checkpoint

This pattern supports both resuming training and loading the best validation model.

Save and Load a Checkpoint

import torch

def save_checkpoint(path, model, optimizer, epoch, metrics, class_names):
    torch.save({
        "epoch": epoch,
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "metrics": metrics,
        "class_names": class_names,
    }, path)

def load_for_training(path, model, optimizer, device):
    checkpoint = torch.load(path, map_location=device)
    model.load_state_dict(checkpoint["model_state_dict"])
    optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
    return checkpoint["epoch"], checkpoint["metrics"]

The model class code must match the saved state dict.
Use map_location when loading on a different device than training.

Inference Function

This function is the shape of a safe inference boundary.

Inference Function

@torch.inference_mode()
def predict(model, batch, device, class_names):
    model.eval()
    batch = batch.to(device)
    logits = model(batch)
    probabilities = torch.softmax(logits, dim=1)
    confidence, class_id = probabilities.max(dim=1)

    return [
        {
            "class": class_names[idx.item()],
            "confidence": conf.item(),
        }
        for conf, idx in zip(confidence, class_id)
    ]

inference_mode is even stricter and faster than no_grad for inference.
Keep preprocessing outside or directly before this function, but make sure it matches training.

Before you move on