Add support for another public dataset like MNIST or CIFAR

[yala]

To help people run the code w.o downloading MIMIC-CXR or CheXpert, add some other public dataset available on torchvision