Fix torchvision pretrained deprecation warning

Spijkervet · May 3, 2024 · 306e762 · 306e762
1 parent f61dde4
commit 306e762
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/simclr/modules/resnet.py b/simclr/modules/resnet.py
@@ -3,8 +3,8 @@
 
 def get_resnet(name, pretrained=False):
     resnets = {
-        "resnet18": torchvision.models.resnet18(pretrained=pretrained),
-        "resnet50": torchvision.models.resnet50(pretrained=pretrained),
+        "resnet18": torchvision.models.resnet18(weights=torchvision.models.ResNet18_Weights.DEFAULT),
+        "resnet50": torchvision.models.resnet50(weights=torchvision.models.ResNet50_Weights.DEFAULT),
     }
     if name not in resnets.keys():
         raise KeyError(f"{name} is not a valid ResNet version")