def get_partition(dataset, start, n_samples):
  """Returns a subset of a given dataset by skipping data.

    Be sure dataset has samples more than `start + n_samples`. Otherwise the
  returned dataset might be an empty one.
  Args:
    dataset: tf.data.Dataset
    start: int, starting index for data.
    n_samples: number_of samples to take from dataset.

  Returns:
    tf.data.Dataset, with `n_samples`.
  """
  return dataset.skip(start).take(n_samples)