fix: load probe class with weights_only = False

adamkarvonen · Feb 14, 2025 · f05bf40 · f05bf40
1 parent 18dc457
commit f05bf40
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/sae_bench/evals/absorption/common.py b/sae_bench/evals/absorption/common.py
@@ -73,6 +73,7 @@ def load_probe(
     probe = torch.load(
         Path(probes_dir) / f"{model_name}" / f"layer_{layer}" / "probe.pth",
         map_location=device,
+        weights_only=False,
     ).to(dtype=dtype)
     return probe