message: Add custom_training_loop test that update keras Metrics.
diff: @@ -988,5 +988,32 @@ class KerasModelsTest(test.TestCase, parameterized.TestCase):
  return model
  +class KerasMetricsTest(test.TestCase, parameterized.TestCase):
  +
  + @combinations.generate(
  + combinations.combine(
  + distribution=strategy_combinations.all_strategies,
  + mode=["eager"]
  + ))
  + def test_multiple_keras_metrics_experimental_run(self, distribution):
  + with distribution.scope():
  + loss_metric = keras.metrics.Mean("loss", dtype=np.float32)
  + loss_metric_2 = keras.metrics.Mean("loss_2", dtype=np.float32)
  +
  + @def_function.function
  + def train_step():
  + def step_fn():
  + loss = constant_op.constant(5.0, dtype=np.float32)
  + loss_metric.update_state(loss)
  + loss_metric_2.update_state(loss)
  +
  + distribution.experimental_run_v2(step_fn)
  +
  + train_step()
  + self.assertEqual(loss_metric.result().numpy(),
  + loss_metric_2.result().numpy())
  + self.assertEqual(loss_metric.result().numpy(), 5.0)
  +
  +
  if __name__ == "__main__":
  test.main()
  