fix for torchvision >= 0.17.2

see also clovaai#215

Author: wkpark

basenet/vgg16_bn.py

@@ -4,7 +4,7 @@
 import torch.nn as nn
 import torch.nn.init as init
 from torchvision import models
-from torchvision.models.vgg import model_urls
+from torchvision.models.vgg import VGG16_BN_Weights
 
 def init_weights(modules):
     for m in modules:
@@ -22,8 +22,10 @@ def init_weights(modules):
 class vgg16_bn(torch.nn.Module):
     def __init__(self, pretrained=True, freeze=True):
         super(vgg16_bn, self).__init__()
-        model_urls['vgg16_bn'] = model_urls['vgg16_bn'].replace('https://', 'http://')
-        vgg_pretrained_features = models.vgg16_bn(pretrained=pretrained).features
+        # Use the weights parameter based on the pretrained flag
+        weights = VGG16_BN_Weights.IMAGENET1K_V1 if pretrained else None
+        vgg_pretrained_features = models.vgg16_bn(weights=weights).features
+
         self.slice1 = torch.nn.Sequential()
         self.slice2 = torch.nn.Sequential()
         self.slice3 = torch.nn.Sequential()