How to condition the sampler

Hi, I've trained the VQVAE and the Pixel snail models on a dataset with multiple classes.

However when I sample, I always seem to be getting different images of the same class.

How can I condition the sampler to give different classes?