fix check in neuron sampling

jbloomAus · Dec 14, 2023 · 809becd · 809becd
1 parent e5f198e
commit 809becd
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/sae_training/sparse_autoencoder.py b/sae_training/sparse_autoencoder.py
@@ -173,7 +173,7 @@ def resample_neurons(
                         print(
                             "Warning: it does not seem as if resetting the Adam parameters worked, there are shapes mismatches"
                         )
-                    if v[v_key][:, replacement_indices].abs().max().item() > 1e-6:
+                    if v[v_key][:, is_dead].abs().max().item() > 1e-6:
                         print(
                             "Warning: it does not seem as if resetting the Adam parameters worked"
                         )