Output only one single class from dataloader

vision
1

Hi,

I’d like to train a VAE on a single digit from the MNIST. Is it possible for the train loader to output only a given class instead of all classes randomly or do I have to write a custom data loader?

Thank you.

2

You can write your own dataset which should inherit from the MNIST dataset.