How do I go about creating a imageFolderDataset that has one sample for each class?
The approach depends a bit on the type of data you are working with.
E.g. if each sample is an image, you could move each image to a separate folder and let ImageFolder create the labels for it.
On the other hand, if you are working with another type of data, you would most likely create the labels manually using some logic, and pass it to TensorDataset or write a custom Dataset to fit your needs.