##
import torch
from torchvision import datasets, transforms, models

def get_device():

    if torch.cuda.is_available():
        # clear cach to boost memory
        # for new round
        torch.cuda.empty_cache()
        return torch.device("cuda")
    else:
        return torch.device("cpu")