fix: Correctly load SAE Bench TopK SAEs (#308)

jbloomAus · Oct 3, 2024 · 4fb5bbe · 4fb5bbe
1 parent 799ce00
commit 4fb5bbe
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/sae_lens/toolkit/pretrained_sae_loaders.py b/sae_lens/toolkit/pretrained_sae_loaders.py
@@ -422,7 +422,7 @@ def get_dictionary_learning_config_1(config: dict[str, Any]) -> dict[str, Any]:
 
     hook_point_name = f"blocks.{trainer['layer']}.hook_resid_post"
 
-    activation_fn_str = "topk" if "topk" in config.get("path", "") else "relu"
+    activation_fn_str = "topk" if trainer["dict_class"] == "AutoEncoderTopK" else "relu"
     activation_fn_kwargs = {"k": trainer["k"]} if activation_fn_str == "topk" else {}
 
     return {