Hi,
I didn’t read your code. But here is what I did to weight the classes such that they are sampled equally (despite the number of samples in each class). First I defined a function called get sampler, it receives: opt (which just has a bool for balance_class), paths which is a list of tuples: containing the full paths of all the data files and the class it belongs to (e…g., paths = [("/data/file1.png", 5), ("/data/file2.png", 6)] … etc ), and class_counts
def get_sampler(opt, paths, class_counts):
    sampler = None
    if opt.balance_class:
        weights = sum(class_counts)/(np.array(class_counts,dtype='float'))
        sampling_weights = []
        for sample in paths:
            sampling_weights.append(weights[sample[1]])
        sampler = data.sampler.WeightedRandomSampler(sampling_weights , len(sampling_weights))
    return sampler

where class_counts  can be computed from  paths
def compute_class_counts(paths):
    class_idxs = [x[1] for x in paths]
    unique, counts = np.unique(class_idxs, return_counts=True)
    return counts

finally, you can pass the output of get_sampler to the Dataloader
sampler = get_sampler(opt, paths, class_counts)
torch.utils.data.DataLoader(
            self.dataset,
            batch_size=opt.batchSize,
            sampler=sampler) # if None does nothing.