Skip to content

Instantly share code, notes, and snippets.

@wcneill
Created July 12, 2020 17:40
Show Gist options
  • Save wcneill/a7674385e592b3ad4a3ff8e4c1aa8af3 to your computer and use it in GitHub Desktop.
Save wcneill/a7674385e592b3ad4a3ff8e4c1aa8af3 to your computer and use it in GitHub Desktop.
# Define pre-trained model
vgg = models.vgg19(pretrained=True).features
# Freeze all model weights
for param in vgg.parameters():
param.requires_grad_(False)
# Move model to GPU, if available.
device = torch.device(
'cuda' if torch.cuda.is_available() else 'cpu'
)
vgg.to(device)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment