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