Solved by crcrpar in post #2 
              
              
                Hi, 
I think the snippet below will do. Note that I’ve not tested this. 

class Dataset(ImageFolder):

    def __init__(self, labels, *args_for_ImageFolder, **kwargs_for_ImageFolder):
        super().__init__(*args_for_ImageFolder, **kwargs_for_ImageFolder)
        self.labels = labels

    def __ge…