Context manager wrapping the training loop, updates step counters.

def training_loop(self):
    """Context manager wrapping the training loop, updates step counters."""
    if not self.restarting:
      self._write_counters(self._local_step_at_start, self._global_step)

    tf.logging.info(
        "Training %s up to %d, %d to go", self.model_mode,
        self.target_local_step, self.steps_to_go
    )

    yield

    self._write_counters(self.target_local_step, -1)