def run(self, n_steps=1):
    """Runs this training loop for n steps.

    Optionally runs evals and saves checkpoints at specified points.

    Args:
      n_steps: Stop training after completing n steps.
    """
    with self._open_summary_writers() as (
        train_summary_writers, eval_summary_writers):
      process = psutil.Process(os.getpid())
      loss_acc, step_acc = 0.0, 0
      start_time = time.time()
      optimizer_metrics_acc = collections.defaultdict(float)
      for i in range(n_steps):
        prev_task_index = self._which_task(self._step)
        self._step += 1
        task_index = self._which_task(self._step)
        task_changed = task_index != prev_task_index

        if task_changed:
          loss_acc, step_acc = 0.0, 0

        loss, optimizer_metrics = self._run_one_step(task_index, task_changed)

        # optimizer_metrics and loss are replicated on self.n_devices, a few
        # metrics are replicated (ex: gradients_l2, weights_l2) - i.e. they are
        # the same across devices, whereas some (ex: loss) aren't because they
        # are different on different devices (due to different data).
        # Taking the average does the correct thing in both the cases.
        #
        # NOTE: Only the weights and gradients are synced across the hosts. This
        # implies the loss here is averaged from this hosts' devices and not
        # across all hosts.
        optimizer_metrics, loss = fastmath.nested_map(
            functools.partial(tl.mean_or_pmean, self._n_devices),
            (optimizer_metrics, loss))

        loss_acc += loss
        # Log loss every 50 steps, every step in memory-efficient trainer.
        if self._step % 50 == 0 or self._use_memory_efficient_trainer:
          self._log_step('Loss: %.4f' % loss, stdout=False)
        step_acc += 1
        for metric_name, value in optimizer_metrics.items():
          optimizer_metrics_acc[metric_name] += value

        # TODO(yuwenyan): Finds a way to log the last round eval step in
        # history.
        #
        # Right now, the last round eval log is missing in history since the
        # checkpoint is saved before it. However sometimes the eval step will
        # fail for some reasons, and it's not acceptable to loose the whole
        # checkpoint in this case. Stays with the old way for now, and fixes it
        # when the checkpoint format is changed to storing weights separately
        # from a small file with history and other data.
        if self._checkpoint_at(self.step):
          self.save_checkpoint('model')
        if self._permanent_checkpoint_at(self.step):
          self.save_checkpoint(f'model_{self.step}')
        if self._eval_at(self.step):
          logging.info('cpu memory use (MB): %.2f',
                       process.memory_info().rss / float(1024 * 1024))
          elapsed_time = time.time() - start_time
          self._log_training_progress(
              task=self._tasks[task_index],
              total_loss=loss_acc,
              n_steps=step_acc,
              elapsed_time=elapsed_time,
              optimizer_metrics=optimizer_metrics_acc,
              summary_writer=train_summary_writers[task_index],
          )
          self.run_evals(eval_summary_writers)
          loss_acc, step_acc = 0.0, 0
          start_time = time.time()
          optimizer_metrics_acc = collections.defaultdict(float)

        # For the current step, after all evals are run and recorded in the
        # event history, check if we need to save special checkpoints because
        # of a new low metric value or a new high metric value.
        if self._checkpoint_at(self.step):
          if self._checkpoint_low_metric is not None and self._at_lowest():
            self.save_checkpoint(f'lowest_{self._checkpoint_low_metric}')
          if self._checkpoint_high_metric is not None and self._at_highest():
            self.save_checkpoint(f'highest_{self._checkpoint_high_metric}')

    # Store the final values back into their respective objects, for testing
    # or other inspection/use.
    #
    # We keep the standard model weights/state unreplicated and
    # tl.Accelerate(model) will carry the replicated weights/state.
    # TODO(afrozm): Try to use tl.Accelerate(model) everywhere in the Loop.
    self._eval_model.weights = self._model.weights
Runs this training loop for n steps.

    Optionally runs evals and saves checkpoints at specified points.

    Args:
      n_steps: Stop training after completing n steps.
    
trax/supervised/training.py