File size: 188 Bytes
d01fde5 |
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 resnet50
return resnet50(weights='IMAGENET1K_V1').to(device)
|