fix broken test

jbloomAus · Apr 10, 2024 · 16935ef · 16935ef
1 parent 1d70af8
commit 16935ef
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/unit/training/test_session_loader.py b/tests/unit/training/test_session_loader.py
@@ -9,6 +9,7 @@
 
 from sae_lens.training.activations_store import ActivationsStore
 from sae_lens.training.config import LanguageModelSAERunnerConfig
+from sae_lens.training.sae_group import SparseAutoencoderDictionary
 
 # from sae_lens.training.sae_group import SAETrainingGroup
 from sae_lens.training.session_loader import LMSparseAutoencoderSessionloader
@@ -152,7 +153,7 @@ def test_load_pretrained_sae_from_huggingface():
         path=folder_path
     )
     assert isinstance(model, HookedTransformer)
-    assert isinstance(sae, SparseAutoencoder)
+    assert isinstance(sae, SparseAutoencoderDictionary)
     assert isinstance(activation_store, ActivationsStore)
     assert sae.cfg.hook_point_layer == layer
     assert sae.cfg.model_name == "gpt2-small"