Merge pull request Sygil-Dev#34 from ZeroCool940711/dev

Changed the pretrained "pretrained=True" to be "weights=torchvision.models.VGG16_Weights.DEFAULT" so it is compatible with latest version of torchvision and also so we always get the latest version of the weights.
ZeroCool940711 · Jun 17, 2023 · 7e88978 · 7e88978
2 parents 85a558f + 6daee00
commit 7e88978
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/muse_maskgit_pytorch/vqgan_vae.py b/muse_maskgit_pytorch/vqgan_vae.py
@@ -407,7 +407,7 @@ def vgg(self):
         if exists(self._vgg):
             return self._vgg
 
-        vgg = torchvision.models.vgg16(pretrained=True)
+        vgg = torchvision.models.vgg16(weights=torchvision.models.VGG16_Weights.DEFAULT)
         vgg.classifier = nn.Sequential(*vgg.classifier[:-2])
         self._vgg = vgg.to(self.device)
         return self._vgg