You could access the .targets attribute from the ImageFolder directly and check the number of unique targets in it via torch.unique(dataset.targets, return_counts=True). Once you know which targets you want to keep, you could filter them out and pass the corresponding sample indices to a Subset, which would then only return the wanted class samples.