B614: pytorch_load

B614: Test for unsafe PyTorch load

This plugin checks for unsafe use of torch.load. Using torch.load with untrusted data can lead to arbitrary code execution. There are two safe alternatives:

Use torch.load with weights_only=True where only tensor data is extracted, and no arbitrary Python objects are deserialized
Use the safetensors library from huggingface, which provides a safe deserialization mechanism

With weights_only=True, PyTorch enforces a strict type check, ensuring that only torch.Tensor objects are loaded.

Example:

>> Issue: Use of unsafe PyTorch load
Severity: Medium   Confidence: High
CWE: CWE-94 (https://cwe.mitre.org/data/definitions/94.html)
Location: examples/pytorch_load_save.py:8
7    loaded_model.load_state_dict(torch.load('model_weights.pth'))
8    another_model.load_state_dict(torch.load('model_weights.pth',
        map_location='cpu'))
9
10   print("Model loaded successfully!")