model.py

import torch
device = "cuda" if torch.cuda.is_available() else "cpu"


def get_model():
    from torchvision.models import alexnet
    return alexnet(weights='IMAGENET1K_V1').to(device)