Using MNIST in pytorch

mnist_torch.py

import torchvision.transforms as transforms
import matplotlib.pyplot as plt

from torch.utils.data import DataLoader
from torchvision.datasets import MNIST

mnist_data = MNIST('./mnist', train=True, download=True, transform=transforms.ToTensor())
data_loader = DataLoader(mnist_data, batch_size=4, shuffle=False)
data_iter = iter(data_loader)

images, labels = data_iter.next()
npimg = images[0].numpy().reshape((28, 28))

plt.imshow(npimg, cmap='gray')
plt.show()