Batches of different distributions

Let’s say that I have a fixed dataset of length N. I want to apply a distribution over this N samples and then sample from this distribution to create mini-batches. I want to continuously change the distribution over my data. One option is to change it after each epoch. This can be easily done by just generating new probabilities for each sample every epoch and then my getitem will just use this distribution to sample an item. However, what I want is to change the distribution after each mini-batch. Is there a way I can do that so that I don’t lose the nice parallelization and multi-threading DataLoader offers? I am not sure where to put the code that generates the probabilities over each sample. Any ideas?