This is based on CIFAR10,
https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html 1
but you can change that to:
trainset=torchvision.dataset.MNIST(**KWARGS)
For filtering out a particular category, you may need to define a custom DataLoader:
https://pytorch.org/tutorials/beginner/basics/data_tutorial.html
and this is good to read, too:
https://pytorch.org/tutorials/beginner/data_loading_tutorial.html 1