[TorchFix] Update deprecated TorchVision pretrained parameters (#1193)

Update deprecated TorchVision pretrained parameters
pytorch · Oct 6, 2023 · c67bbab · c67bbab
1 parent cead596
commit c67bbab
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/cpp/transfer-learning/convert.py b/cpp/transfer-learning/convert.py
@@ -5,7 +5,7 @@
 from torchvision import models
 
 # Download and load the pre-trained model
-model = models.resnet18(pretrained=True)
+model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
 
 # Set upgrading the gradients to False
 for param in model.parameters():

diff --git a/fast_neural_style/neural_style/vgg.py b/fast_neural_style/neural_style/vgg.py
@@ -7,7 +7,7 @@
 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=models.VGG16_Weights.IMAGENET1K_V1).features
         self.slice1 = torch.nn.Sequential()
         self.slice2 = torch.nn.Sequential()
         self.slice3 = torch.nn.Sequential()