From 24434a7eabe97c17595c5754c6f9e02a40696bc2 Mon Sep 17 00:00:00 2001 From: tailingy <186656915+tailingy@users.noreply.github.com> Date: Wed, 25 Mar 2026 18:13:16 +0000 Subject: [PATCH] fix(polygraphy): use weights_only=True in torch.load to prevent RCE via pickle deserialization Passing weights_only=True to torch.load() restricts deserialization to tensor data only, blocking arbitrary pickle-based code execution. Fixes NVBugs #5934574 - [PSIRT] Arbitrary Code Execution via unsafe torch.load() in Polygraphy JSON deserialization. --- tools/Polygraphy/polygraphy/json/serde.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/Polygraphy/polygraphy/json/serde.py b/tools/Polygraphy/polygraphy/json/serde.py index 607457b83..9c29bd90a 100644 --- a/tools/Polygraphy/polygraphy/json/serde.py +++ b/tools/Polygraphy/polygraphy/json/serde.py @@ -265,7 +265,7 @@ def encode(tensor): def decode(dct): data = base64.b64decode(dct["tensor"].encode(), validate=True) infile = io.BytesIO(data) - return torch.load(infile) + return torch.load(infile, weights_only=True) TORCH_REGISTRATION_SUCCESS = True