diff --git a/align_anything/trainers/text_to_text/kto.py b/align_anything/trainers/text_to_text/kto.py
index d484886a..bbf0a3d6 100644
--- a/align_anything/trainers/text_to_text/kto.py
+++ b/align_anything/trainers/text_to_text/kto.py
@@ -67,6 +67,8 @@ def compute_kl(self):
             sampler=DistributedSampler(random_dataset, shuffle=True),
             batch_size=self.cfgs.train_cfgs.per_device_kl_batch_size,
         )
+        kl_sum = 0
+        num_batches = 0
         for batch in self.random_dataloader:
             log_probs = self.compute_log_probs(  # size = (2 * B, L - 1)
                 self.model.module,
@@ -78,7 +80,10 @@ def compute_kl(self):
             )
             kl = (log_probs - ref_log_probs).mean()
 
-            self.kl = max(kl, 0)
+            kl_sum = kl_sum + max(kl, 0)
+            num_batches += 1
+
+        self.kl = kl_sum / num_batches
 
     def loss(  # pylint: disable=too-many-locals
         self,