adding topk to hidden pre test

jbloomAus · Nov 12, 2024 · 85524bc · 85524bc
1 parent 2185e52
commit 85524bc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/unit/training/test_training_sae.py b/tests/unit/training/test_training_sae.py
@@ -181,7 +181,7 @@ def test_TrainingSAE_forward_includes_topk_loss_is_nonzero_if_dead_neurons_prese
     assert train_step_output.losses["auxiliary_reconstruction_loss"] > 0.0
 
 
-@pytest.mark.parametrize("architecture", ["standard", "gated", "jumprelu"])
+@pytest.mark.parametrize("architecture", ["standard", "gated", "jumprelu", "topk"])
 def test_TrainingSAE_encode_returns_same_value_as_encode_with_hidden_pre(
     architecture: str,
 ):