fix: fixing type errors after bad merge

jbloomAus · Nov 10, 2024 · 4a08d0d · 4a08d0d
1 parent ceb2d3f
commit 4a08d0d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/unit/training/test_training_sae.py b/tests/unit/training/test_training_sae.py
@@ -29,11 +29,11 @@ def test_TrainingSAE_training_forward_pass_can_scale_sparsity_penalty_by_decoder
 
     if scale_sparsity_penalty_by_decoder_norm:
         assert (
-            pytest.approx(train_step_output.l1_loss)
+            pytest.approx(train_step_output.losses["l1_loss"].detach().item())  # type: ignore
             == 2.0 * scaled_feature_acts.norm(p=1, dim=1).mean().detach().item()
         )
     else:
         assert (
-            pytest.approx(train_step_output.l1_loss)
+            pytest.approx(train_step_output.losses["l1_loss"].detach().item())  # type: ignore
             == 2.0 * feature_acts.norm(p=1, dim=1).mean().detach().item()
         )