# Uncomment and run the commands below if imports fail
# !conda install numpy pandas pytorch torchvision cpuonly -c pytorch -y
# !pip install matplotlib --upgrade --quiet

import torch
import torchvision
import numpy as np
import matplotlib.pyplot as plt
import torch.nn as nn
import torch.nn.functional as F
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor
from torchvision.utils import make_grid
from torch.utils.data.dataloader import DataLoader
from torch.utils.data import random_split
%matplotlib inline

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

VAL_SIZE=10000
TRAIN_SIZE=len(dataset)-VAL_SIZE

train_ds, val_ds=random_split(dataset, [TRAIN_SIZE,VAL_SIZE])
len(train_ds),len(val_ds)

#Data loader
BATCH_SIZE=128
train_loader=DataLoader(train_ds,BATCH_SIZE,shuffle=True, num_workers=4,pin_memory=True)
val_loader=DataLoader(val_ds,BATCH_SIZE*2,shuffle=True, num_workers=4,pin_memory=True)