If you want a predefined number of samples from each class, you could implement a custom sampler, which would yield the batch indices using your desired logic.
On the other hand, if your use case would allow weighted random sampling with replacement, you could use a WeightedRandomSampler as described here.