def evaluate(self, weights):
    """Evaluates the model and returns the metrics."""
    self._predict_model.weights = weights

    metrics = collections.defaultdict(list)
    for _ in range(self._n_steps):
      batch = self._eval_task.next_batch()
      step_metrics = self._evaluate_batch(batch)
      for (key, value) in step_metrics.items():
        metrics[key].append(value)

    metrics = {k: np.array(v) for (k, v) in metrics.items()}

    def metric_name(context, horizon):
      return f'pred_error/context_{context}/horizon_{horizon}'

    return {
        metric_name(context, horizon):
            np.sum(errors) / (np.sum(errors != 0) + 1e-6)
        for ((context, horizon), errors) in metrics.items()
    }
Evaluates the model and returns the metrics.
trax/supervised/callbacks.py