import cv2 as cv
import os
import torch.nn as nn
import torch
import torchvision
import tarfile
import torch.functional as f
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torch.optim
from torchvision.datasets.utils import download_url
from torch.utils.data import DataLoader as DataLoader
from google.colab.patches import cv2_imshow
from zipfile import ZipFile
import glob
from google.colab import files
import imageio
from albumentations import HorizontalFlip, VerticalFlip, Rotate
import matplotlib.pyplot as plt

from google.colab import drive
drive.mount('/content/drive')

train_ds  =  datasets.MNIST(root  = '/data', train=True, transform= transforms.ToTensor(), download=True)
test_ds  =  datasets.MNIST(root  = '/data', train=False, transform= transforms.ToTensor(), download=True)
print(len(train_ds))
print(len(test_ds))

batch_size  = 64
input_size  = 784
output      = 10
lr          = 0.001
num_epoch   = 2
loss_function = nn.CrossEntropyLoss()