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.
     """