How to randomly choose a subset of classes every batch and sample from those chosen classes?

Hi all,

I want to get samples from a random subset of classes and sample a batch of examples for this iteration. A running example would be, for training Cifar100, for each iteration in a single epoch, I want to get a minibatch of 256 images but consisting only of randomly chosen 50 of 100 classes w/o replacement, so each class would have around 5 examples. And for the next minibatch, I want another different 50 classes.

I am thinking of using a custom sampler to do this and get indices of examples of the same classes and then sample from them, but it would be costly.

Any thoughts would be appreciated!