Internal change

PiperOrigin-RevId: 320536301

official/nlp/tasks/masked_lm.py

@@ -62,10 +62,10 @@ class MaskedLMTask(base_task.Task):
       sentence_labels = labels['next_sentence_labels']
       sentence_outputs = tf.cast(
           model_outputs['next_sentence'], dtype=tf.float32)
-      sentence_loss = tf.keras.losses.sparse_categorical_crossentropy(
-          sentence_labels,
-          sentence_outputs,
-          from_logits=True)
+      sentence_loss = tf.reduce_mean(
+          tf.keras.losses.sparse_categorical_crossentropy(sentence_labels,
+                                                          sentence_outputs,
+                                                          from_logits=True))
       metrics['next_sentence_loss'].update_state(sentence_loss)
       total_loss = mlm_loss + sentence_loss
     else:

official/nlp/tasks/sentence_prediction.py

@@ -91,7 +91,7 @@ class SentencePredictionTask(base_task.Task):
 
     if aux_losses:
       loss += tf.add_n(aux_losses)
-    return loss
+    return tf.reduce_mean(loss)
 
   def build_inputs(self, params, input_context=None):
     """Returns tf.data.Dataset for sentence_prediction task."""