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.