org.sonar.l10n.py.rules.python.S6985.html Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of python-checks Show documentation

There is a newer version: 4.23.0.17664

This rule raises an issue when pytorch.load is used to load a model.
Why is this an issue?
In PyTorch, it is common to load serialized models using the torch.load function. Under the hood, torch.load uses the
pickle library to load the model and the weights. If the model comes from an untrusted source, an attacker could inject a malicious
payload which would be executed during the deserialization.
How to fix it
Use a safer alternative to load the model, such as safetensors.torch.load_model. Alternatively, PyTorch can be instructed to only load
the weights by setting the parameter weights_only=True. This avoids the use of the pickle library and is therefore safe.
Note that the use of weights_only requires saving only the state_dict of a model instead of the whole model.
Code examples
Noncompliant code example
import torch

model = torch.load('model.pth') # Noncompliant: torch.load is used to load the model

Compliant solution
import torch
import safetensors

model = MyModel()
safetensors.torch.load_model(model, 'model.pth')

Resources
Documentation

   Pytorch documentation: Save/Load Entire
  Model