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)
|
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)
|