fix accidental bug

jbloomAus · Apr 21, 2024 · c22fbbd · c22fbbd
1 parent b5e97f8
commit c22fbbd
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/sae_lens/training/sparse_autoencoder.py b/sae_lens/training/sparse_autoencoder.py
@@ -312,13 +312,12 @@ def load_from_pretrained_legacy(cls, path: str):
 
         # Create an instance of the class using the loaded configuration
         instance = cls(cfg=state_dict["cfg"])
-        new_state_dict = instance.state_dict()
         if "scaling_factor" not in state_dict["state_dict"]:
             assert isinstance(instance.cfg.d_sae, int)
             state_dict["state_dict"]["scaling_factor"] = torch.ones(
                 instance.cfg.d_sae, dtype=instance.cfg.dtype, device=instance.cfg.device
             )
-        instance.load_state_dict(new_state_dict, strict=True)
+        instance.load_state_dict(state_dict["state_dict"], strict=True)
 
         return instance
 

diff --git a/tests/unit/toolkit/test_pretrained_saes.py b/tests/unit/toolkit/test_pretrained_saes.py
@@ -24,7 +24,7 @@ def test_convert_old_to_modern_saelens_format():
 
     # convert file format
     pretrained_saes.convert_old_to_modern_saelens_format(
-        legacy_out_file, new_out_folder
+        legacy_out_file, new_out_folder, force=True
     )
 
     # Load from new converted file