Solved by avalon1511 in post #11 
              
              
                This is how I’m sampling equally from each class of the dataset 
def _create_samples(dataset, num_classes):

        N = int(np.ceil(k_samp / num_classes)) # k_samp is the number of total samples I need
        indices = np.arange(len(dataset))
        train_indices, test_indices = train_test_split(…