Internal change

PiperOrigin-RevId: 366819679

official/core/base_task.py

@@ -217,8 +217,14 @@ class Task(tf.Module, metaclass=abc.ABCMeta):
     with tf.GradientTape() as tape:
       outputs = model(features, training=True)
       # Computes per-replica loss.
-      loss = self.build_losses(
-          labels=labels, model_outputs=outputs, aux_losses=model.losses)
+      if model.compiled_loss:
+        loss = model.compiled_loss(
+            labels, outputs, regularization_losses=model.losses)
+        loss += self.build_losses(
+            labels=labels, model_outputs=outputs, aux_losses=None)
+      else:
+        loss = self.build_losses(
+            labels=labels, model_outputs=outputs, aux_losses=model.losses)
       # Scales loss as the default gradients allreduce performs sum inside the
       # optimizer.
       scaled_loss = loss / tf.distribute.get_strategy().num_replicas_in_sync

official/core/base_trainer_test.py

@@ -368,6 +368,18 @@ class TrainerTest(tf.test.TestCase, parameterized.TestCase):
     with self.assertRaises(RuntimeError):
       _ = trainer.train(tf.convert_to_tensor(2, dtype=tf.int32))
 
+  def test_model_with_compiled_loss(self):
+    task = mock_task.MockTask()
+    model = task.build_model()
+    model.compile(loss=tf.keras.losses.CategoricalCrossentropy())
+    trainer = trainer_lib.Trainer(
+        self._config,
+        task,
+        model=model,
+        optimizer=task.create_optimizer(self._config.trainer.optimizer_config))
+    logs = trainer.train(tf.convert_to_tensor(5, dtype=tf.int32))
+    self.assertIn('training_loss', logs)
+
 
 if __name__ == '__main__':
   tf.test.main()