From 36dab85319464218500963bd7159e6ace6cbf230 Mon Sep 17 00:00:00 2001
From: Geaming <71267655+Geaming2002@users.noreply.github.com>
Date: Sat, 25 Jan 2025 02:27:42 +0800
Subject: [PATCH] chore: fix sae loading in logits_lens_with_features.ipynb
 example (#406)

when load `sae` through `SAE.from_pretrained`, `sae_id` should change with `layer` not always `blocks.0.hook_resid_pre`
---
 tutorials/logits_lens_with_features.ipynb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tutorials/logits_lens_with_features.ipynb b/tutorials/logits_lens_with_features.ipynb
index 69edd924..b319cec4 100644
--- a/tutorials/logits_lens_with_features.ipynb
+++ b/tutorials/logits_lens_with_features.ipynb
@@ -130,7 +130,7 @@
     "for layer in range(12):\n",
     "    sae, original_cfg_dict, sparsity = SAE.from_pretrained(\n",
     "        release=\"gpt2-small-res-jb\",\n",
-    "        sae_id=\"blocks.0.hook_resid_pre\",\n",
+    "        sae_id=f\"blocks.{layer}.hook_resid_pre\",\n",
     "        device=\"cpu\",\n",
     "    )\n",
     "    gpt2_small_sparse_autoencoders[f\"blocks.{layer}.hook_resid_pre\"] = sae\n",