What is the recommended strategy for loading legacy PyTorch checkpoints after the torch.load weights_only default changed to True?

Recommended Strategy

The recommended strategy for loading trusted legacy checkpoints containing custom Python objects is to use torch.serialization.add_safe_globals() to whitelist the specific classes required for reconstruction. This allows you to maintain the security benefits of the weights_only=True default while restoring the ability to load complex state dictionaries.

Implementation Steps

1. Identify Required Classes: Attempt to load the checkpoint with the default weights_only=True. If the checkpoint contains custom objects, PyTorch will raise a RuntimeError explicitly naming the classes that are not in the safe list.
2. Whitelist the Globals: Use the add_safe_globals API to register these classes before calling torch.load():

   import torch
   from my_module import CustomOptimizerState
   
   torch.serialization.add_safe_globals([CustomOptimizerState])
   checkpoint = torch.load("legacy_model.pt")
   

3. Verify State Integrity: Compare the keys and types of the loaded state_dict against a known-good legacy load (using weights_only=False in a controlled environment) to ensure no metadata was silently dropped.

Technical Considerations

Nested Containers: You do not need to whitelist standard Python containers like list or dict; the restricted unpickler handles these natively. You only need to whitelist the leaf custom types contained within those structures.

Resaving vs. Whitelisting: While add_safe_globals solves the immediate loading issue, the most robust long-term strategy is to resave trusted checkpoints using a weights-only compatible format, such as safetensors. This removes the dependency on pickle entirely and eliminates the need for version-specific allowlists.

Security and Versioning

Caution: Never set weights_only=False for checkpoints sourced from untrusted third parties, as this exposes the system to arbitrary code execution via pickle.

Assumptions: This guidance assumes the use of PyTorch 2.x versions where the weights_only parameter and add_safe_globals API are available. If you are using an older version, these parameters will not exist, and torch.load will behave as weights_only=False by default.

Diagnostic Detail Needed: To refine this recommendation, please specify if your checkpoints contain torch.nn.Module subclasses saved as entire objects (via torch.save(model)) rather than state dictionaries (via torch.save(model.state_dict())), as the latter is significantly easier to secure.