fix(polygraphy): use weights_only=True in torch.load to prevent RCE

[tailingy]

Summary

• Vulnerability: Arbitrary Code Execution (RCE) via unsafe pickle deserialization in polygraphy/json/serde.py
• Root Cause: torch.load(infile) was called without weights_only=True, allowing maliciously crafted Polygraphy JSON files to execute arbitrary code via Python's pickle module
• Fix: Pass weights_only=True to torch.load(), restricting deserialization to tensor data only

Change

# Before (vulnerable)
return torch.load(infile)

# After (safe)
return torch.load(infile, weights_only=True)

Impact

An attacker who can supply a crafted Polygraphy JSON file (e.g. via the CLI or Python API) can execute arbitrary code in the victim's process context. This enables data exfiltration, privilege escalation, and other attacks.

Compatibility

weights_only=True has been supported since PyTorch 1.13, which is already the minimum version declared by Polygraphy (torch>=1.13.0). No compatibility issues expected.

References

• NVBugs #5934574 - [PSIRT] Tensor-RT LLM pickle deserialization pytorch
• PyTorch docs: torch.load

🤖 Generated with Claude Code