diff --git a/crosscode/trainers/topk_crosscoder/trainer.py b/crosscode/trainers/topk_crosscoder/trainer.py
index b5824b7..336f66f 100644
--- a/crosscode/trainers/topk_crosscoder/trainer.py
+++ b/crosscode/trainers/topk_crosscoder/trainer.py
@@ -36,9 +36,9 @@ def _calculate_loss_and_log(
                 **self._get_fvu_dict(batch_BMPD, train_res.recon_acts_BMPD),
             }
 
-            return reconstruction_loss, log_dict
+            return loss, log_dict
 
-        return reconstruction_loss, None
+        return loss, None
 
     def aux_loss(
         self, batch_BMPD: torch.Tensor, train_res: ModelHookpointAcausalCrosscoder.ForwardResult