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(indices, train_size = N * num_classes , stratify = dataset.targets)

        # Warp into Subsets and DataLoaders
        
        train_dataset = torch.utils.data.Subset(target_dataset, train_indices)

        return train_dataset