Upload folder using huggingface_hub

steerable_retrieval/steer/loading.py

@@ -111,7 +111,10 @@ def load_steerable_sae(
     if not hasattr(dec, "b_dec"):
         dec.b_dec = enc.b_dec
 
-    state = torch.load(ckpt_path, map_location=device)
+    # Lightning checkpoints from our training runs can contain OmegaConf metadata
+    # alongside tensors. PyTorch 2.6 defaults torch.load(weights_only=True), which
+    # rejects that metadata; this loader is for trusted project checkpoints.
+    state = torch.load(ckpt_path, map_location=device, weights_only=False)
     sd = state.get("state_dict", state)
     enc_sd = {k[len("sae_encoder."):]: v for k, v in sd.items() if k.startswith("sae_encoder.")}
     dec_sd = {k[len("sae_decoder."):]: v for k, v in sd.items() if k.startswith("sae_decoder.")}