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)