[GKD] Fix batchmean reduce op in GKDTrainer's loss (#4105)

Author: cmpatino

trl/trainer/gkd_trainer.py

@@ -270,7 +270,7 @@ def generalized_jsd_loss(
 
         # Apply reduction
         if reduction == "batchmean":
-            return jsd.sum() / mask.sum() if labels is not None else jsd.sum() / (jsd.size(0) * jsd.size(1))
+            return jsd.sum() / mask.sum() if labels is not None else jsd.sum() / jsd.size(0)
         elif reduction == "sum":
             return jsd.sum()
         elif reduction == "mean":