def __init__(
      self,
      loop,
      model=None,
      eval_at=1000,
      eval_task=None,
      context_lengths=(1,),
      horizon_lengths=(1,),
      n_steps=1,
      accelerate_model=True,
  ):
    """Initializes SerializedModelEvaluation.

    Args:
      loop: Instance of `trax.supervised.training.Loop` or `None`. Can be set to
        `None` for testing - in such a case, `model` and `eval_task` must be
        provided.
      model: Instance of `trax.rl.serialization_utils.SerializedModel`. Not
        required if `loop` is provided.
      eval_at: When to evaluate. Either int (every how many steps to evaluate),
        or a list of ints (step numbers), or a function int -> bool (step
        predicate).
      eval_task: Instance of `trax.supervised.training.EvalTask` with the
        evaluation data, or None. If not provided, the task will be taken from
        `loop`.
      context_lengths: List of lengths of the context sequence fed into the
        model before starting prediction.
      horizon_lengths: List of lengths of the predicted sequence.
      n_steps: Number of batches to run evaluation for.
      accelerate_model (bool): Whether to wrap the model in `tl.Accelerate`.
    """
    super().__init__(loop)

    if model is None:
      model = loop.model

    observation_serializer = model.observation_serializer
    action_serializer = model.action_serializer

    predict_model = model.make_predict_model()
    if accelerate_model:
      predict_model = tl.Accelerate(predict_model)
    self._predict_model = predict_model
    self._obs_serializer = observation_serializer
    self._act_serializer = action_serializer

    if isinstance(eval_at, int):
      self._eval_at = lambda step: step % eval_at == 1
    elif hasattr(eval_at, '__in__'):
      self._eval_at = lambda step: step in eval_at
    elif callable(eval_at):
      self._eval_at = eval_at
    else:
      raise TypeError(f'Unsupported type for eval_at: {type(eval_at)}.')

    if eval_task is None:
      if len(loop.eval_tasks) != 1:
        raise ValueError(
            'If eval_task is not provided, the number of eval_tasks registered '
            'in Loop must be exactly 1.'
        )
      eval_task = loop.eval_tasks[0]
    self._eval_task = eval_task

    self._context_lengths = list(sorted(context_lengths))
    self._horizon_lengths = list(sorted(horizon_lengths))
    self._n_steps = n_steps

    self._batch_size = eval_task.sample_batch[0].shape[0]
    (_, self._init_state) = predict_model.init(
        shapes.ShapeDtype((self._batch_size, 1), dtype=np.int32)
    )
Initializes SerializedModelEvaluation.

    Args:
      loop: Instance of `trax.supervised.training.Loop` or `None`. Can be set to
        `None` for testing - in such a case, `model` and `eval_task` must be
        provided.
      model: Instance of `trax.rl.serialization_utils.SerializedModel`. Not
        required if `loop` is provided.
      eval_at: When to evaluate. Either int (every how many steps to evaluate),
        or a list of ints (step numbers), or a function int -> bool (step
        predicate).
      eval_task: Instance of `trax.supervised.training.EvalTask` with the
        evaluation data, or None. If not provided, the task will be taken from
        `loop`.
      context_lengths: List of lengths of the context sequence fed into the
        model before starting prediction.
      horizon_lengths: List of lengths of the predicted sequence.
      n_steps: Number of batches to run evaluation for.
      accelerate_model (bool): Whether to wrap the model in `tl.Accelerate`.
    
trax/supervised/callbacks.py