Introduction

Saving and loading models in PyTorch is the process of persisting a trained model's parameters (and optionally its training state) to disk, and later restoring them — without needing to retrain from scratch. PyTorch's approach differs meaningfully from the TensorFlow/Keras saving covered earlier in this curriculum: rather than saving a complete, self-contained model file by default, PyTorch's standard and recommended approach saves only the model's state_dict — its learned parameters — which must then be loaded back into an already-instantiated model whose class definition already exists in your code.

Understanding this distinction, along with when and why you'd reach for TorchScript or ONNX Export (covered in the previous two topics) instead, completes the full picture of how a trained PyTorch model actually gets persisted, restored, and prepared for its next stage of use.

Why Does Saving and Loading Matter?

Saving and loading models helps to:

  • Avoid the time and compute cost of retraining a model every time it's needed
  • Enable resuming training after an interruption, without losing prior progress
  • Support checkpointing during long training runs (as covered in the Validation and Callbacks topics)
  • Allow a trained model to be shared, reused, or fine-tuned later
  • Preserve not just model weights, but optionally optimizer state for exact training resumption
  • Provide the foundation that TorchScript and ONNX Export build on when preparing a model for deployment

What Gets Saved: state_dict

A state_dict is simply a Python dictionary mapping each layer's
name to its corresponding learned parameter tensor. This is the
same automatic parameter tracking covered in the nn.Module
topic — every weight and bias registered as part of the model
is captured here, in a clean, structured format.

Saving and Loading state_dict (The Recommended Approach)

Notice that loading requires first creating a new instance of
the SAME model class before the weights can be loaded into it —
PyTorch's official documentation specifically recommends this
approach over saving the entire model object, since it's more
stable across different PyTorch versions and doesn't tie the
saved file to Python's pickle serialization of the exact class
definition.

Saving the Entire Model (Less Recommended)

While this appears more convenient, PyTorch's official
documentation actively recommends AGAINST this approach for
anything beyond quick, throwaway scripts. It relies on Python's
pickle module to serialize the exact class definition alongside
the data, which can break if the code defining that class moves,
changes, or isn't available in the exact same form when loading
— making state_dict the safer, more portable, recommended
standard for real projects.

Saving a Full Training Checkpoint (For Resuming Training)

This connects directly to the Optimizers topic: an optimizer
like Adam maintains its own internal state (moving averages of
past gradients), and saving optimizer.state_dict() alongside
the model's weights allows training to resume EXACTLY where it
left off, rather than restarting the optimizer's internal
tracking from scratch — important for genuinely continuing
a long, interrupted training run correctly.

state_dict Only vs Full Model Save

Aspectstate_dict (Recommended)Full Model Object
What's SavedJust the learned parametersThe entire model object, including its class structure
Requires Original Class Code?Yes — must instantiate the same class before loadingIn theory, no — but fragile in practice
Portability/StabilityHigh — recommended, more robust across PyTorch versionsLower — can break if code changes or class isn't available
File SizeSmallerLarger
PyTorch's Official RecommendationYesNo — recommended against for real projects

Saving state_dict vs Using TorchScript/ONNX

Aspectstate_dictTorchScriptONNX Export
Requires Python + Original Class?YesNo (once converted)No
Primary PurposeContinuing work within PyTorch/PythonDeployment via LibTorch, mobile, C++Cross-framework, cross-runtime deployment
Captures Architecture?No — only weightsYes — architecture and weights togetherYes — architecture and weights together
Best ForEveryday development, resuming training, fine-tuningPython-independent PyTorch deploymentFramework-agnostic deployment

Moving a Loaded Model to a Device

The map_location argument is particularly important when a model was saved on a GPU but needs to be loaded on a machine without one (or vice versa) — it tells PyTorch explicitly which device to map the saved tensors to during loading, avoiding device-mismatch errors.

Key Properties of Saving and Loading in PyTorch

  • PyTorch's recommended approach saves only a model's state_dict, not the entire model object.
  • Loading a state_dict requires first instantiating the same model class already defined in your code.
  • Saving a full training checkpoint (model + optimizer state dicts) enables exact resumption of interrupted training.
  • Saving the entire model object works but is officially discouraged due to fragility across code and version changes.
  • map_location should be used when loading a model onto a different device than the one it was saved from.

Where Does Saving and Loading Matter Most?

ContextWhy It Matters
Long Training RunsCheckpointing (model + optimizer state) protects against interruptions
Fine-Tuning Pre-Trained ModelsLoading existing weights as a starting point for further training
Model SharingDistributing trained weights for others to load into a matching model class
Experimentation WorkflowsSaving multiple model versions to compare or roll back to later
Production PreparationThe starting point before conversion to TorchScript or ONNX for deployment

Advantages

  • state_dict's simple dictionary format is stable, portable, and easy to inspect
  • Saving optimizer state alongside model weights enables exact training resumption
  • Straightforward, well-documented pattern used consistently throughout the PyTorch ecosystem
  • Small, efficient file sizes compared to saving an entire model object
  • map_location cleanly handles moving saved models across different devices

Limitations

  • Requires the original model class definition to be available when loading a state_dict
  • Saving the full model object (the less-recommended path) can silently break with code changes
  • Doesn't provide the Python-independence that TorchScript or ONNX Export offer
  • Manual checkpoint management (tracking epochs, best models, etc.) must be implemented yourself, as covered in the Callbacks topic
  • Version mismatches between PyTorch versions can occasionally cause loading compatibility issues

Real-World Examples

ApplicationSaving/Loading Use
Long-Running Model TrainingFull checkpoints (model + optimizer) to survive interruptions
Transfer Learning ProjectsLoading pre-trained weights as a starting point for fine-tuning
Model Sharing Between ResearchersDistributing state_dict files alongside the model class code
A/B Testing Different Model VersionsSaving and comparing multiple trained state_dict checkpoints
Production Deployment PipelinesLoading a trained state_dict before converting to TorchScript or ONNX

Best Practices

  • Use state_dict saving and loading as the default approach, following PyTorch's official recommendation.
  • Save full checkpoints (model + optimizer state) for any training run you might need to resume later.
  • Always call model.eval() after loading a model intended for inference, not further training.
  • Use map_location explicitly when loading a model onto a different device than it was saved from.
  • Keep the original model class code alongside saved state_dict files, since it's required to load them back correctly.

Interview Tip

A common interview question is:

"Why does PyTorch recommend saving a model's state_dict instead of saving the entire model object?"

A strong answer is:

PyTorch recommends saving just the state_dict — a dictionary of the model's learned parameters — because saving the entire model object relies on Python's pickle module to serialize the exact class definition alongside the data, which can break if the code defining that class changes, moves, or isn't available in exactly the same form later, especially across different PyTorch versions. Saving only the state_dict and requiring the model class to already be instantiated when loading is more stable, portable, and considered the safer standard, even though it means you need access to the original model-defining code to load the weights back correctly.

Explaining the pickle-related fragility gives a concrete, technical reason that makes your answer stronger.

Conclusion

Saving and loading with state_dict provides PyTorch's standard, recommended way to persist and restore trained models, checkpoints, and optimizer state within the Python/PyTorch ecosystem — complementing the TorchScript and ONNX Export paths covered previously for scenarios that require moving beyond Python entirely. With Installation, Tensors, Autograd, Dataset, DataLoader, nn.Module, Optimizers, Training Loop, Validation, Callbacks, TorchScript, ONNX Export, and now Save/Load Models all covered, this completes the full PyTorch section, spanning the entire journey from initial setup through training, evaluation, and deployment preparation.