import torch
import torchvision
from torchvision.datasets import MNIST

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

len(dataset)

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

import matplotlib.pyplot as plt
%matplotlib inline

image, label = dataset[10]
plt.imshow(image)
print('Label: ',label)

import torchvision.transforms as transforms
dataset = MNIST(root='data/', train = True, transform = transforms.ToTensor())
img_tensor, label = dataset[0]
print(img_tensor.shape,label)