Solved by techkang in post #2 
              
              
                You can implement both (i) and (ii) like this: 
import random

from torch.utils.data import DataLoader, Sampler, RandomSampler, Dataset


class MyBatchSampler(Sampler):
    def __init__(self, sampler, batch_size_list, drop_last):
        self.sampler = sampler
        self.batch_size_list = batch_si…