fix input args

2025-05-20 07:34:41 -07:00
parent 6fafe46562
commit 83ad248e5b
1 changed files with 1 additions and 1 deletions
--- 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,