a2e95a60dc10751706d84ea2656b9f72b7b18110,official/utils/misc/keras_utils.py,ProfilerCallback,on_batch_end,#ProfilerCallback#Any#Any#,151

Before Change

  def on_batch_end(self, batch, logs=None):
    if batch == self.stop_step_in_epoch and self.should_stop:
      self.should_stop = False
      profiler.stop()
      logging.info("Profiler saved profiles for steps between %s and %s to %s",
                   self.start_step, self.stop_step, self.log_dir)

After Change

  def on_batch_end(self, batch, logs=None):
    if batch == self.stop_step_in_epoch and self.should_stop:
      self.should_stop = False
      results = profiler.stop()
      profiler.save(self.log_dir, results)
      logging.info(
          "Profiler saved profiles for steps between %s and %s to %s",
          self.start_step, self.stop_step, self.log_dir)

Instances