Hi, I am trying to manually define training batches for each epoch such that each batch contains only samples that satisfying specific condition (e.g. first batch are all from class 1 with image size smaller than 128*128, second batch are all from class 2 with image size greater than 256*256 and so on, and batches may have different number of samples).  That said, using for batch in dataloader would probably not give me correct grouping of batches.  I am wondering if anyone know any good way to manually “assign batches” before each training epoch starts?
Thank you!