Update examples/fast_neural_style/vgg.py

pytorch · Jul 6, 2023 · d52e94b · d52e94b
1 parent 35a2da3
commit d52e94b
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/examples/fast_neural_style/vgg.py b/examples/fast_neural_style/vgg.py
@@ -2,12 +2,12 @@
 
 import torch
 from torchvision import models
-
+from torchvision.models.vgg16 import VGG16_Weights
 
 class Vgg16(torch.nn.Module):
     def __init__(self, requires_grad=False):
         super(Vgg16, self).__init__()
-        vgg_pretrained_features = models.vgg16(pretrained=True).features
+        vgg_pretrained_features = models.vgg16(weights=VGG16_Weights.IMAGENET1K_V1).features
         self.slice1 = torch.nn.Sequential()
         self.slice2 = torch.nn.Sequential()
         self.slice3 = torch.nn.Sequential()
@@ -36,3 +36,4 @@ def forward(self, X):
         vgg_outputs = namedtuple("VggOutputs", ["relu1_2", "relu2_2", "relu3_3", "relu4_3"])
         out = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3)
         return out
+