'Get the idx\'th batch from the dataset.'

def get_batch(self, idx):

    assert (idx >= 0), 'idx    must    be    non    negative'
    assert (idx < self.num_batches), 'idx    must    be    less    than    the    number    of    batches'
    start_idx = (idx * self.batch_size)
    indices = range(start_idx, (start_idx + self.batch_size))
    return self._get_batch_from_indices(indices)
