Save and load a PyTorch model: checkpoints that resume training

To save a PyTorch model for inference, save its state_dict and reload it into the matching architecture. To resume training, save the training state as well: the optimizer, learning-rate scheduler, progress counter, and any state that determines the next batch or update.

A working checkpoint lets you recover from an interrupted training job without repeating the whole run. A successful file write is only the first check: start a new process, restore the checkpoint, and compare its next training step with an uninterrupted run. This guide provides a small CPU example that does exactly that, then explains what changes for GPU and distributed jobs.

Choose between model weights and a training checkpoint

PyTorch’s saving and loading guide recommends saving a model’s state dictionary rather than serializing the whole model object. A checkpoint dictionary can also hold the optimizer and progress information needed to continue a run.

Keep the architecture, configuration, preprocessing, and dependency versions with the saved files. Weights cannot tell your application how to tokenize a sentence, resize an image, or interpret a prediction. If you fine-tune an adapter, also record the base model and its revision; our LoRA fine-tuning guide covers that workflow.

• Inference: restore the trained model parameters and persistent buffers, then configure evaluation.
• Resume training: restore the model and the training loop’s relevant state.
• Start a different experiment: loading weights with a new optimizer may be intentional, but record it as a new run.

Load a checkpoint file with explicit settings

Use torch.load(path, map_location="cpu", weights_only=True) for the plain tensor-and-dictionary files in this guide. The loading API lets you choose where tensor storage is restored; CPU mapping is useful when the saving machine had a GPU that the loading machine does not.

The weights_only argument restricts deserialization; it does not mean the file can contain only model weights. The example includes optimizer state, tensors, dictionaries, numbers, and strings. PyTorch’s serialization notes explain the supported loading behavior and its security limits. Use files from sources you trust. Do not switch to weights_only=False just to silence an unexpected loading error.

A tested save-and-resume training example

This example uses a fixed synthetic batch, a small neural network with dropout, AdamW, and a step-based learning-rate schedule. It saves after four completed updates, then computes a fifth update as the reference. A second process restores step four and checks that its fifth update matches.

Save the complete code below as pytorch_checkpoint_demo.py. Run it in a separate working directory with PyTorch installed. It creates a checkpoint-demo folder and overwrites only its two named demonstration checkpoint files when run again. The exact example was tested with Python 3.12.14 and PyTorch 2.15.0+cpu on Windows 11, using one CPU thread. No GPU is required.

import argparse
from pathlib import Path
import torch

CONFIG = {"inputs": 4, "hidden": 8, "outputs": 1, "dropout": 0.25}


def build(config):
    model = torch.nn.Sequential(
        torch.nn.Linear(config["inputs"], config["hidden"]),
        torch.nn.ReLU(),
        torch.nn.Dropout(config["dropout"]),
        torch.nn.Linear(config["hidden"], config["outputs"]),
    )
    optimizer = torch.optim.AdamW(model.parameters(), lr=0.01)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=2, gamma=0.5)
    return model, optimizer, scheduler


def train_step(model, optimizer, scheduler):
    # Fixed synthetic batch; no data-loader state is needed for this example.
    x = torch.arange(16, dtype=torch.float32).reshape(4, 4) / 10
    y = x.sum(dim=1, keepdim=True)
    model.train()
    optimizer.zero_grad(set_to_none=True)
    loss = torch.nn.functional.mse_loss(model(x), y)
    loss.backward()
    optimizer.step()
    scheduler.step()  # This example schedules once per optimizer update.
    return loss.item()


def snapshot(model, optimizer, scheduler, step, loss):
    return {
        "format_version": 1,
        "torch_version": str(torch.__version__),
        "model_config": CONFIG,
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "scheduler": scheduler.state_dict(),
        "completed_steps": step,
        "last_loss": loss,
        "torch_rng": torch.get_rng_state(),
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("mode", choices=["save", "resume"])
    parser.add_argument("--directory", default="checkpoint-demo")
    args = parser.parse_args()
    folder = Path(args.directory)
    folder.mkdir(parents=True, exist_ok=True)
    torch.set_num_threads(1)
    torch.use_deterministic_algorithms(True)

    if args.mode == "save":
        torch.manual_seed(7)
        model, optimizer, scheduler = build(CONFIG)
        for step in range(1, 5):
            loss = train_step(model, optimizer, scheduler)
        path = folder / "step-4.pt"
        torch.save(snapshot(model, optimizer, scheduler, step, loss), path)
        loss = train_step(model, optimizer, scheduler)
        torch.save(snapshot(model, optimizer, scheduler, 5, loss),
                   folder / "reference-step-5.pt")
        print(f"Saved {path}; reference step 5 loss: {loss:.8f}")
        return

    checkpoint = torch.load(folder / "step-4.pt", map_location="cpu", weights_only=True)
    if checkpoint["format_version"] != 1:
        raise ValueError("Unsupported checkpoint format")
    model, optimizer, scheduler = build(checkpoint["model_config"])
    model.load_state_dict(checkpoint["model"])
    scheduler.load_state_dict(checkpoint["scheduler"])
    optimizer.load_state_dict(checkpoint["optimizer"])
    # Restore RNG after constructing objects, immediately before the next step.
    torch.set_rng_state(checkpoint["torch_rng"])
    loss = train_step(model, optimizer, scheduler)
    step = checkpoint["completed_steps"] + 1
    reference = torch.load(folder / "reference-step-5.pt", map_location="cpu", weights_only=True)
    assert step == reference["completed_steps"]
    assert loss == reference["last_loss"]
    for key, value in model.state_dict().items():
        torch.testing.assert_close(value, reference["model"][key], rtol=0, atol=0)
    assert scheduler.state_dict() == reference["scheduler"]
    assert optimizer.param_groups[0]["lr"] == reference["optimizer"]["param_groups"][0]["lr"]
    print(f"Resume matched reference step {step}; loss: {loss:.8f}")


if __name__ == "__main__":
    main()

Run the save and resume commands separately

python pytorch_checkpoint_demo.py save
python pytorch_checkpoint_demo.py resume

In the recorded test, the resumed process printed Resume matched reference step 5; loss: 10.83563042. The script checks the step counter, loss, model tensors, scheduler state, and learning rate. The numeric loss is specific to this toy run; it is not a model-quality or hardware-performance result.

Additional checks compared the full optimizer state and tested incomplete restores. Omitting optimizer state changed the updated model. Omitting the random-number state changed the loss and model. Omitting scheduler state left its counter wrong even though this particular next update still matched. That last result is why one matching loss is insufficient evidence of a correct resume.

Restore the training state in the right order

Construct the model, optimizer, and scheduler before loading their saved states. Keep the optimizer’s parameter groups consistent with the original run. PyTorch’s optimizer loading documentation warns that initializing a scheduler after loading the optimizer can overwrite restored learning rates. The example initializes both first, loads the scheduler state, then loads the optimizer state.

Restore the random-number state after object construction, immediately before continuing. Constructing model layers can consume random numbers. In this example, the restored CPU generator controls dropout; the input batch itself is fixed.

Define the counter precisely. Here, completed_steps=4 means four optimizer updates finished, so the next update is step five. In an epoch-based loop, distinguish the last completed epoch from the next epoch to run. With gradient accumulation, prefer a checkpoint boundary after a completed optimizer update unless your implementation also preserves the unfinished accumulation state.

Save and load model weights for inference

Using the same model-building function, the following example extracts inference weights from the training checkpoint and loads them into a new model. Keep it beside pytorch_checkpoint_demo.py and run it after the save command.

from pathlib import Path
import torch
from pytorch_checkpoint_demo import build

folder = Path("checkpoint-demo")
checkpoint = torch.load(folder / "step-4.pt", map_location="cpu", weights_only=True)
model, _, _ = build(checkpoint["model_config"])
model.load_state_dict(checkpoint["model"])
torch.save(model.state_dict(), folder / "inference-weights.pt")

loaded_model, _, _ = build(checkpoint["model_config"])
weights = torch.load(folder / "inference-weights.pt", map_location="cpu", weights_only=True)
loaded_model.load_state_dict(weights)
loaded_model.eval()
with torch.inference_mode():
    prediction = loaded_model(torch.ones(1, 4))
print(prediction)

model.eval() selects evaluation behavior for layers such as dropout. torch.inference_mode() controls gradient-related work; it does not call eval() for you. The inference-mode documentation makes this distinction explicit. Return to model.train() when continuing training.

Keep the best model and the latest checkpoint separately

A model selected by validation loss and the last checkpoint answer different questions. The best model is the candidate you want to evaluate or deploy; the latest checkpoint limits how much training you repeat after an interruption. Keep both, with the metric name, step, and run identifier in their records.

Do not keep an in-memory “best” snapshot using only best = model.state_dict() while training continues. The Module API specifies that the returned dictionary contains references to parameters and buffers. Save it immediately or make an independent copy. Otherwise, later updates can change the state you thought you preserved.

What changes for GPU, mixed precision, and shuffled data?

Reproducibility has a defined scope

Record the hardware, software versions, data order, and precision settings used for a resume test. PyTorch does not guarantee identical results across releases, platforms, or CPU and GPU execution; see its reproducibility guidance. Saving one seed is not the same as saving every generator’s current state.

A real input pipeline may need Python and NumPy random state, CUDA generator states, sampler state, data-loader generator state, and a position within the dataset. Worker processes, random augmentation, and prefetching make a mid-epoch resume more involved. Choose a boundary your data pipeline can actually restore and test the next batch as well as the next update.

Restore mixed-precision state when the run uses it

If training uses a gradient scaler, save and restore its state_dict alongside the model and optimizer. The official AMP recipe describes saving at an iteration boundary, including after scaler.update(). This CPU example does not test AMP or CUDA. For installation and device checks, use our PyTorch on A100 guide.

Store checkpoints where a replacement job can read them

A checkpoint on a disappearing instance’s local disk cannot help a replacement job. Choose storage with the required persistence and access, then test restoration from a separate process or machine. For large shared workloads, the HPC file-systems guide explains the storage choices.

Choose a save interval around the amount of work you can afford to repeat and the measured cost of writing a checkpoint. Keep versioned files and at least one known-good earlier checkpoint. Finish writing and verifying a new file before treating it as the latest usable copy; on remote storage, confirm the upload completed. Avoid deleting the previous recovery point merely because a new filename exists.

The demonstration uses synchronous local writes and does not implement a production storage policy. Production code also needs failure handling, retention rules, capacity monitoring, and a way to associate each file with its model code and dataset revision.

Use the right checkpointing method for larger jobs

Activation checkpointing and durable training checkpoints solve different problems. Activation checkpointing recomputes intermediate results during backward to reduce memory use. It does not save a recovery file for the next machine.

For sharded models, use a checkpoint path designed for your distributed training stack. PyTorch Distributed Checkpoint handles state spread across ranks and has its own compatibility requirements. Do not assume one process’s ordinary state_dict file contains the whole distributed model.

Common save and load questions

Why not save the entire model object?

torch.save(model, path) couples the file to the model’s Python class definitions and import paths. Refactoring that code can break loading. A state dictionary makes the stored parameters explicit, though you still need compatible model code to reconstruct the architecture.

Does .pt versus .pth change what is saved?

Those are naming conventions. Inspect the saved structure: a model state dictionary and a training checkpoint dictionary require different loading code. Use a format version in your own checkpoint schema so incompatible changes fail clearly.

Should I use strict=False to fix a loading error?

First check the architecture, parameter shapes, key names, and file type. A partial load can be useful for an intentional transfer-learning experiment, but it is not a substitute for restoring the expected model. Inspect any missing or unexpected keys instead of ignoring them.

Should new deployment work start with TorchScript?

TorchScript is deprecated; PyTorch’s 2.14 release notes retain that status. Choose a deployment route appropriate to the target runtime, such as a supported torch.export workflow. Exporting an inference artifact is a separate task from preserving resumable training state.