message: Assert scalar summary correctly written with Keras and Custom Training Loop.
diff: @@ -246,10 +246,10 @@ class AutoOutsideCompilationWithKerasTest(test.TestCase):
  def testSummaryWithCustomTrainingLoop(self):
  strategy = get_tpu_strategy()
  + writer = summary_ops_v2.create_file_writer_v2(self.summary_dir)
  with strategy.scope():
  model = distribute_strategy_test.get_model()
  model.compile('sgd', 'mse')
  - writer = summary_ops_v2.create_file_writer_v2(self.summary_dir)
  @def_function.function
  def custom_function(dataset):
  @@ -273,6 +273,15 @@ class AutoOutsideCompilationWithKerasTest(test.TestCase):
  distribute_strategy_test.get_dataset(strategy))
  custom_function(dataset)
  + writer.close()
  +
  + event_files = file_io.get_matching_files_v2(
  + os.path.join(self.summary_dir, 'event*'))
  + events_count_dictionary = {
  + ('logits'): 0,
  + }
  + self.validate_recorded_sumary_file(event_files, events_count_dictionary,
  + 1)
  if __name__ == '__main__':
  