How to use BatchSampler to choose from different indexes in different batches

Hi,

I have a custom dataset, that is pretty straightforward (path, image, label). For each label, I know it relevant indexes. For the training phase, I need each batch to sample from one label only (my architecture supports that). For example - dusring the first 3 batches sample only from ‘label1’, during the 4th to 6th batch sample from ‘label2’ only.
It looks like BatchSampler, that wraps another sampler to yield mini-batches, is the way to go, but I’m not sure how it could support that in order to create one dataloder.

Assuming I have a custom dataset - “train_dataset”, and I know the indexes of all the targets, such as:
indx_dict = {‘label1’: [20,21,22,23,24,25],
‘label2’ : [25,26,27,28,29, 30]}

How can I create one dataloader that samples twice for example from “label1” indexes (bs=3), then samples twice from “label2” indexes (bs=3)?
Preferably do that with replacement.
The solution should support more than only 2 labels, that’s just a short example