import torch
from torchvision import datasets
import torchvision.transforms as tf

train_dataset = datasets.MNIST('./data', train=True, download=True, transform=tf.ToTensor())
test_dataset = datasets.MNIST('./data', train=False, download=True, transform=tf.ToTensor())