update deprecated VGG parameters

tg-93 · Sep 5, 2022 · 5a345c1 · 5a345c1
1 parent 651835a
commit 5a345c1
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/basicsr/archs/vgg_arch.py b/basicsr/archs/vgg_arch.py
@@ -101,11 +101,11 @@ def __init__(self,
                 max_idx = idx
 
         if os.path.exists(VGG_PRETRAIN_PATH):
-            vgg_net = getattr(vgg, vgg_type)(pretrained=False)
+            vgg_net = getattr(vgg, vgg_type)(weights=None)
             state_dict = torch.load(VGG_PRETRAIN_PATH, map_location=lambda storage, loc: storage)
             vgg_net.load_state_dict(state_dict)
         else:
-            vgg_net = getattr(vgg, vgg_type)(pretrained=True)
+            vgg_net = getattr(vgg, vgg_type)(weights=vgg.VGG19_Weights.DEFAULT)
 
         features = vgg_net.features[:max_idx + 1]