def METHOD_NAME(self, b, n):
    indices = self.get_codebook_indices()
    indices = indices.view(-1, self.groups)
    cb_size = indices.size(0)
    assert (
        n < cb_size
    ), f"sample size {n} is greater than size of codebook {cb_size}"
    sample_idx = torch.randint(low=0, high=cb_size, size=(b * n,))
    indices = indices[sample_idx]
    z = self.vars.squeeze(0).index_select(0, indices.flatten()).view(b, n, -1)
    return z