import torch
from torch.utils.data import DataLoader, random_split
from torch.nn.functional import cross_entropy, mse_loss
import torch.nn.functional as F
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from torchvision.utils import make_grid
from torchvision.datasets import ImageFolder
from torchvision.datasets.utils import download_url
import matplotlib.pyplot as plt
from tqdm.notebook import tqdm
import tarfile
import random
import numpy as np

# Dowload the dataset
dataset_url = "http://files.fast.ai/data/cifar10.tgz"
download_url(dataset_url, '.')

# Extract from archive
with tarfile.open('./cifar10.tgz', 'r:gz') as tar:
    tar.extractall(path='./data')

dataset = ImageFolder('./data/cifar10/train/', transform=transforms.ToTensor())
test_set = ImageFolder('./data/cifar10/test/', transform=transforms.ToTensor())

def get_default_device():
    """Pick GPU if available, else CPU"""
    if torch.cuda.is_available():
        return torch.device('cuda')
    else:
        return torch.device('cpu')
device = get_default_device()
device