model.py

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


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