Hello!

Today's topic: just because an attractive PyTorch checkpoint is being handed around does not mean you should innocently use it — doing so can be seriously dangerous.

Have you ever thought about the security risks of loading a model checkpoint? A checkpoint you assumed was just a model weights file can turn out to be an unexpected security hole.

In this article, we look at the safety of PyTorch's torch.load function and present practical guidelines for handling model checkpoints properly.

The Hidden Danger in Model Checkpoints

PyTorch's torch.load function is extremely convenient, but it carries a serious security risk.

Here is why.

• A checkpoint is not just parameters!
  A checkpoint file can contain not only the model's weights but also arbitrary Python code — classes, functions, and more.
  
• It can contain executable code!
  This means a checkpoint is not merely a "data file" — it is a file with the potential to execute Python code.
  
• Which means there is a potential vulnerability!
  If you carelessly load a malicious checkpoint, arbitrary code can be executed the moment the file is loaded.

That's right — a checkpoint is not just weight data.

What Can Go into a Checkpoint File

So, a PyTorch checkpoint file (.pt, .pth, or .ckpt) can end up containing not only the model's weights (parameters) but actual Python code as well.

This is because PyTorch internally uses Python's pickle serialization format.

pickle is a mechanism for saving and restoring the state of Python objects, but by its nature it can also store executable code, including class definitions and functions.

Concretely, a checkpoint may contain elements such as the following.

• Model weight parameters
• Model architecture information
• Optimizer state
• Learning-rate scheduler state
• Custom class and function definitions(← this is the dangerous part)
• Other metadata

The "custom class and function definitions" are the security concern. A checkpoint crafted by a malicious actor can contain dangerous code — deleting files, executing system commands, and so on.

That is why innocently using an attractive checkpoint just because someone is handing it out can be seriously dangerous.

A Concrete Attack Scenario

The problems a malicious checkpoint can cause unfold like this

1. (You) download a malicious file
2. (You) load it with weights_only=False.
3. The malicious code embedded in the checkpoint (planted by the attacker) executes immediately
4. (In the worst case) the attacker can perform arbitrary operations on your PC — deleting files, exfiltrating data, installing malware, and more

Security Hardening in PyTorch 2.6

PyTorch 2.6 and later strengthened the security measures. The specific change is that

• torch.load()'s weights_only argument now defaults to True.
• As a result, by default only the model's parameters (weights) are loaded safely, preventing execution of potentially dangerous code

Because of this change, you may encounter errors such as the following

RuntimeError: ('Attempted to deserialize object from torch.nn.Module that contains non-parameter/buffer types, which could potentially lead to security vulnerabilities. ...

 _pickle.UnpicklingError: Weights only load failed. This file can still be loaded, to do so you have two options, do those steps only if you trust the source of the checkpoint.
         (1) In PyTorch 2.6, we changed the default value of the weights_only argument in torch.load from False to True. Re-running torch.load with weights_only set to False will likely succeed, but it can result in arbitrary code execution. Do it only if you got the file from a trusted source.
         (2) Alternatively, to load with weights_only=True please check the recommended steps in the following error message.
         WeightsUnpickler error: Unsupported global:...

This error occurs when the checkpoint file contains not just model weights but also code information such as classes and functions. For example, if a class like fairseq.data.dictionary.Dictionary is included, the default settings reject it as an unapproved class.

Safety Criteria: What Counts as "Safe" and "Unsafe"

🟢 Files that can be considered "safe"

The following cases can be regarded as essentially "safe"

• Checkpoints of models you created yourself(because you fully understand their contents)
• Model files obtained directly from trusted official sources (well-known universities, reputable companies, official repositories, and so on)
• Models obtained from official GitHub repositories or official websites, where the provider is widely known and the contents of the model checkpoint are explicitly documented (e.g., official Hugging Face repositories, Facebook's official fairseq repository)

(Even so, you still need to keep a careful eye on security advisories, and loading with weights_only=True remains the default rule.)

🔴 Files that may be "unsafe"

You should suspect the following cases as "unsafe":

• Models from sites of unknown origin or shared anonymously by third parties
• Model files downloaded directly from unidentifiable individuals on social media, unofficial forums, and the like
• Cases where the distributor of the model file cannot be identified and insufficient information about the model's structure and code has been published

You must neverweights_only=False load such files with it.

🔐 Code Examples and Remedies for Safe Model Loading

How to resolve the error when it occurs

Here we look at an error that actually occurred while using Facebook's fairseq, and how to deal with it

 _pickle.UnpicklingError: Weights only load failed. This file can still be loaded, to do so you have two options, do those steps only if you trust the source of the checkpoint.
         (1) In PyTorch 2.6, we changed the default value of the weights_only argument in torch.load from False to True. Re-running torch.load with weights_only set to False will likely succeed, but it can result in arbitrary code execution. Do it only if you got the file from a trusted source.
         (2) Alternatively, to load with weights_only=True please check the recommended steps in the following error message.
         WeightsUnpickler error: Unsupported global: GLOBAL fairseq.data.dictionary.Dictionary was not an allowed global by default. Please use torch.serialization.add_safe_globals([Dictionary]) or the torch.serialization.safe_globals([Dictionary]) context manager to allowlist this global if you trust this class/function.

This error occurred because PyTorch's torch.load() now applies stricter default settings for security when loading a model checkpoint.

Specifically, it comes down to the following

• In PyTorch 2.6 and later, torch.load()'s weights_only argument defaults to True, so it attempts to safely load only the model's parameters (weights).
• The checkpoint file that triggers this error contains not just model weights but also code information such as classes and functions. Therefore, weights_only=True attempting to load it with the default setting is rejected on the grounds that it contains unapproved classes or functions.
• In this particular error, we can read from the message that the class rejected as unapproved is fairseq.data.dictionary.Dictionary.

(1) Recommended solution (high safety)

Register the trusted class on an "allowlist" before loading.

import torch
from fairseq.data.dictionary import Dictionary

# Run this only if you know this class is safe
torch.serialization.add_safe_globals([Dictionary])

# Then load the checkpoint
checkpoint = torch.load('checkpoint.pt')

Alternatively, to allow it only temporarily, use the context manager.

import torch
from fairseq.data.dictionary import Dictionary

with torch.serialization.safe_globals([Dictionary]):
    checkpoint = torch.load('checkpoint.pt')

(2) Legacy approach (low safety)

weights_only=False can be specified when loading. However, this approach is not safe, so use it only for files obtained from trusted sources.

import torch

checkpoint = torch.load('checkpoint.pt', weights_only=False)

Implementation Example: A Function That Loads Checkpoints Safely

Below is an example of a checkpoint-loading function designed with safety in mind

import torch
from typing import List, Type, Optional

def load_checkpoint_safely(
    checkpoint_path: str,
    trusted_classes: Optional[List[Type]] = None,
    force_weights_only: bool = False
) -> dict:
    """
    Load a checkpoint safely.
    
    Args:
        checkpoint_path: Path to the checkpoint file
        trusted_classes: List of trusted classes (empty list if None)
        force_weights_only: If True, force loading with weights_only=True
        
    Returns:
        The contents of the checkpoint
        
    Raises:
        RuntimeError: If loading fails
    """
    if force_weights_only:
        return torch.load(checkpoint_path, weights_only=True)
        
    trusted_classes = trusted_classes or []
    
    try:
        # First, try to load safely
        return torch.load(checkpoint_path)
    except RuntimeError as e:
        if "non-parameter/buffer types" in str(e):
            # Use the trusted classes if any are provided
            if trusted_classes:
                with torch.serialization.safe_globals(trusted_classes):
                    return torch.load(checkpoint_path)
            else:
                raise RuntimeError(
                    f"The checkpoint contains unapproved classes. "
                    f"If it comes from a trusted source, use the trusted_classes parameter. "
                    f"Original error: {str(e)}"
                )
        else:
            # Re-raise any other errors as-is
            raise

Finally, let's wrap up with the best practices drawn from this case

Best Practices Summary

1. Default to safety first
   • Take advantage of the default setting in PyTorch 2.6 and later (weights_only=True)
   • For files of unknown trustworthiness, always attempt to load with the default settings first
     
2. Use an explicit allowlist
   • Explicitly allow trusted classes with add_safe_globals or the safe_globals context manager
   • This lets you allow only the minimum necessary classes and preserve safety
     
3. weights_only=False is a last resort
   • Use it only for trusted files, and only as a last resort
   • Never use it on checkpoints from unknown sources

Conclusion: Raising Our Security Awareness

The security hardening in PyTorch 2.6 and later is a step forward for the safe sharing and use of deep learning models.

The basic safety criterion is very simple

• Clear provenance and trustworthy (with evidence of safety) → safe
• Unknown origin, anonymous, or suspicious provider → unsafe

When we moved our own products to PyTorch 2.6, code that had been working stopped working, and quite a few revisions were needed. But thanks to that, we feel our awareness and knowledge of model security took a real step forward.

As deep learning applications and players multiply, many useful models are now being distributed — but more players also means more players with bad intentions, so we intend to keep exercising great care when loading checkpoints. If you need to develop products on older versions of PyTorch, you should be especially careful.

See you next time!