From 3a0faa97cabecc08d7b87c10b387e2e86bd5e31b Mon Sep 17 00:00:00 2001
From: Wing Lian <wing@axolotl.ai>
Date: Wed, 28 May 2025 00:46:52 -0400
Subject: [PATCH] fix trainer callback base class

---
 src/axolotl/integrations/kd/callbacks.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/axolotl/integrations/kd/callbacks.py b/src/axolotl/integrations/kd/callbacks.py
index b8a806f69..521833477 100644
--- a/src/axolotl/integrations/kd/callbacks.py
+++ b/src/axolotl/integrations/kd/callbacks.py
@@ -4,10 +4,10 @@ Transformers trainer callbacks to schedule the KD temperature during training
 
 import math
 
-from transformers.trainer_callback import CallbackHandler
+from transformers.trainer_callback import TrainerCallback
 
 
-class KDTemperatureSchedulerCallback(CallbackHandler):
+class KDTemperatureSchedulerCallback(TrainerCallback):
     """
     KD temperature scheduler callback for the trainer.
     """