import torch 
import torchvision
from torchvision.datasets import MNIST

# download the dataset
#dataset = MNIST(root = 'data/', download = True)
dataset = MNIST(root = 'data/')

len(dataset)

test_dataset = MNIST(root = 'data/', train = False)

len(test_dataset)