diff --git a/src/xorl/server/runner/model_runner.py b/src/xorl/server/runner/model_runner.py
index e5b28236..48ac29ea 100644
--- a/src/xorl/server/runner/model_runner.py
+++ b/src/xorl/server/runner/model_runner.py
@@ -1904,6 +1904,7 @@ def _compute_micro_batch_loss(self, micro_batch, loss_fn, loss_fn_params):
                 ce_mode=self.ce_mode,
                 compute_kl_stats=compute_kl_stats,
                 lm_head_fp32=self.lm_head_fp32,
+                loss_reducer=token_sum_reducer,
                 metric_reducer=token_sum_reducer,
             )
             local_loss_sum = _result.loss
@@ -2010,6 +2011,7 @@ def _compute_micro_batch_loss(self, micro_batch, loss_fn, loss_fn_params):
                 compute_kl_stats=compute_kl_stats,
                 lm_head_fp32=self.lm_head_fp32,
                 icepop_beta=icepop_beta,
+                loss_reducer=token_sum_reducer,
                 metric_reducer=token_sum_reducer,
             )
             local_loss_sum = _result.loss