diff --git a/src/axolotl/integrations/kd/trainer.py b/src/axolotl/integrations/kd/trainer.py
index 0e2996504..b688252d3 100644
--- a/src/axolotl/integrations/kd/trainer.py
+++ b/src/axolotl/integrations/kd/trainer.py
@@ -92,7 +92,7 @@ class AxolotlKDTrainer(AxolotlTrainer):
             )
         else:
             loss_kd = self.kd_loss_fn(
-                # shift_logits,
+                shift_logits,
                 target_token_ids_for_loss,
                 target_logprobs_for_loss,
                 target_mask_for_loss,