def create_init_batch(self, features):
    """Returns a batch of size "hparams.init_batch_size" for initialization.

    Args:
      features: input features.
    Returns:
      init_features: initialization features.
    """
    train_dataset = self.hparams.problem.dataset(
        tf_estimator.ModeKeys.TRAIN, hparams=self.hparams)
    train_dataset = train_dataset.batch(self.hparams.init_batch_size)
    train_dataset = self.init_preprocess(train_dataset)
    return train_dataset.make_one_shot_iterator().get_next()
Returns a batch of size "hparams.init_batch_size" for initialization.

    Args:
      features: input features.
    Returns:
      init_features: initialization features.
    
tensor2tensor/models/research/glow.py