File size: 186 Bytes
f30fd8c |
1 2 3 4 5 6 7 8 |
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)
|