def _filter_cols(dataset, seed):
  """Mapping function.

  Filter columns for batch.
  Args:
    dataset: tf.data.Dataset with several columns.
    seed: int
  Returns:
    dataset: tf.data.Dataset with two columns (X, Y).
  """
  col_y = f'image/sim_{seed}_y/value'
  return dataset['image/encoded'], dataset[col_y]