Solved by Mika_Spilner in post #4 
              
              
                Finally, I made a custom BatchSampler. 
I hope this code help someone like me. 
Reference here 
import numpy as np
import torch

from torch.utils.data import DataLoader
from torch.utils.data.sampler import BatchSampler


class BalancedBatchSampler(BatchSampler):
    """
    BatchSampler - from a MNI…