import torch
import torchvision
import numpy as np
from torchvision.datasets import MNIST
from torchvision import transforms

data_transforms = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean = (0.1, ), std = (0.1, )),
])

dataset = MNIST(
    root = "data/",
    train=True,
    transform = data_transforms,
    download=True,
)

len(dataset)

from torch.utils.data import random_split