def train_epoch(self, n_steps, n_eval_steps):
    """Runs `n_steps` of training, with periodic logging, saving, and evals."""
    # TODO(jonni): Clarify how this method relates to the stricter notion of
    # epoch (training for as many steps as needed for a full pass through the
    # training data).
    print()  # Add visual separator in logs for start of training epoch.
    start_time = time.time()

    for _ in range(n_steps):
      batch = next(self._train_stream)
      if self.n_devices > 1:  # TODO(lukaszkaiser): use everywhere if possible.
        batch = _reshape_by_device(batch, self.n_devices)
      self.train_step(batch)
      if self._should_save_now():
        self.save_state(keep=True)
      if self._should_log_now():
        self._train_sw.scalar('training/learning_rate', self.learning_rate)

    # At end of n_steps, do bookkeeping, run evals, and save state.
    elapsed_time = time.time() - start_time
    self.log_step('Ran %d train steps in %0.2f secs' % (n_steps, elapsed_time))
    if self._train_sw and n_steps > 1:
      self._train_sw.scalar('training/steps per second',
                            n_steps / elapsed_time, step=self._step)
      self._train_sw.flush()
    self.evaluate(n_eval_steps)
    if self._eval_sw:
      self._eval_sw.flush()
    if self._should_save_checkpoints:
      self.save_state(keep=False)
    if self._should_save_checkpoints and self._current_step_is_best(high=True):
      self.save_state(keep=False, prefix='highest_' + self._checkpoint_highest)
    if self._should_save_checkpoints and self._current_step_is_best(high=False):
      self.save_state(keep=False, prefix='lowest_' + self._checkpoint_lowest)
Runs `n_steps` of training, with periodic logging, saving, and evals.
trax/supervised/trainer_lib.py