From c0ca06c07b6fa8568d335db98fb1c1787507bf08 Mon Sep 17 00:00:00 2001 From: Eldruin Date: Thu, 2 Jul 2026 00:10:45 -0300 Subject: [PATCH] fix: use weights_only=True in torch.load to prevent unsafe deserialization --- examples/sampling_pytorch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/sampling_pytorch.py b/examples/sampling_pytorch.py index dc156c2..03dce52 100644 --- a/examples/sampling_pytorch.py +++ b/examples/sampling_pytorch.py @@ -67,7 +67,7 @@ def _load_and_sample( """Loads and samples a string from a checkpoint.""" device = "cuda" if torch.cuda.is_available() else "cpu" print(f"Loading the parameters from {path_checkpoint}") - params = torch.load(path_checkpoint) + params = torch.load(path_checkpoint, weights_only=True) params = {k: v.to(device=device) for k, v in params.items()} print("Parameters loaded.") # Create a sampler with the right param shapes.