From 83ad248e5bf17a368ef03c525c9cb49fd12d0325 Mon Sep 17 00:00:00 2001
From: Wing Lian <wing@axolotl.ai>
Date: Tue, 20 May 2025 07:34:41 -0700
Subject: [PATCH] fix input args

---
 src/axolotl/integrations/kd/trainer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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,