Load VGG model and pytorch essentials

st_blog_1

if torch.cuda.is_available():
    device = torch.device('cuda')
else:
    raise Exception('GPU is not available')
    
# Load VGG19 features. We do not need the last linear layers,
# only CNN layers are needed
vgg = vgg19(pretrained=True).features
vgg = vgg.to(device)
# We don't want to train VGG
for param in vgg.parameters():
    param.requires_grad_(False)
    
torch.backends.cudnn.benchmark = True