use kd_alpha in the correct loss method

2024-12-24 19:54:32 -05:00
parent 3416302b0d
commit ca5e397fc5
1 changed files with 2 additions and 1 deletions
--- a/src/axolotl/core/trainers/kd.py
+++ b/src/axolotl/core/trainers/kd.py
@@ -182,7 +182,8 @@ class AxolotlKDTrainer(AxolotlTrainer):
        )

        if self.args.kd_ce_alpha > 0:
-            loss = self.args.kd_ce_alpha * outputs["loss"] + loss_kd
+            kd_alpha = self.args.kd_alpha
+            loss = self.args.kd_ce_alpha * outputs["loss"] + kd_alpha * loss_kd
        else:
            loss = loss_kd
        # Save past state if it exists