def _batch_dataset(self, dataset):
    """Batch the dataset depending on what model is used.

    Args:
      dataset: A dataset to batch

    Returns:
      A batched dataset with correct padding shapes.
    """
    if self._encoder is None:
      raise ValueError('Cannot batch dataset without an encoder.')
    return dataset.padded_batch(
        batch_size=self._model_config.batch_size,
        padded_shapes=(
            self._encoder.padded_shapes(),
            # Must match losses in training.py
            {
                'claim_classification': [],
                'evidence_matching': []
            }))
Batch the dataset depending on what model is used.

    Args:
      dataset: A dataset to batch

    Returns:
      A batched dataset with correct padding shapes.
    
language/serene/training.py