def make_collater_fn(
      config: ml_collections.ConfigDict
  ) -> Callable[[Dict[str, tf.Tensor]], Dict[str, tf.Tensor]]:
    """Produces function to preprocess batches. See BaseTask."""

    def collater_fn(batch: Dict[str, tf.Tensor]) -> Dict[str, tf.Tensor]:
      return batch

    return collater_fn
Produces function to preprocess batches. See BaseTask.
language/mentionmemory/tasks/example_task.py