I wrote a function that returns a balanced sampler for SubsetRandomSampler, which can be used as as a sampler in Dataloder(s). The function is working well and might be useful for others. However, I wonder if anyone has any comments, suggestions or improvements.
import numpy as np
import torch

def get_a_balanced_sampler(score, sampler_size, labels, no_classes): 
    
    ''' 
    Args in - 
        score: posteriori values in the range [0 to 1] for each of the labels, 
        a value 1 indicates that the label has high likeliness to be of a correct 
        class
        sampler_size: the intended sampler size the user wants to get back, the 
        size of the returned sampler will be (slightly) less than this, depending on 
        the minimum number of labels per class
        labels: an array containing the labels
        no_classes: the number of classes in the problem         
        
    Parameters - 
        percentage_of_selected_samples: selecting 50% for the samples with the highest 
        'score' values, thus, selection will be made randomly from these samples in a 
        balanced manner. The 50% can be changed according to the user requirements, 
        e.g., using less or higher values.
    '''
    percentage_of_selected_samples = 50/100 
    
    len_labels_per_class = np.zeros(no_classes, dtype=int)      
    idx_per_class = np.zeros([no_classes, len(labels)], dtype=int)
    for i in range(no_classes):
        idx_per_class[i] = labels==i
        len_labels_per_class[i] = sum(idx_per_class[i] == True)
    no_labels_per_class = min(len_labels_per_class)   
    sampler_pool_size = int(no_labels_per_class * percentage_of_selected_samples) 
    
    sampler_size = int(sampler_size/no_classes)
    if(sampler_size > sampler_pool_size): 
        print('You need to decrease the value percentage_of_selected_samples: ', percentage_of_selected_samples)
        exit('Exiting functget_the_samplerpler(): sampler_size has become larger than sampler_pool_size')
            
    my_sampler = []
    for i in range(no_classes):
        sample_idx = (-score[idx_per_class[i]]).argsort()[:sampler_pool_size]   
        sample_idx = np.random.permutation(sample_idx)
        sample_idx = sample_idx[:sampler_size]
        my_sampler.extend(sample_idx)
    
    if len(my_sampler) < 100:  exit('Exiting function get_a_balanced_sampler(): small sampler has been geneated')    
    my_sampler = torch.utils.data.sampler.SubsetRandomSampler(my_sampler)  
    
    return my_sampler