Fix deprecated parameters in test for pretrained weights

ayrna · Dec 20, 2023 · e684f8c · e684f8c
1 parent 65d6fab
commit e684f8c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/dlordinal/estimator/tests/test_estimator.py b/dlordinal/estimator/tests/test_estimator.py
@@ -10,7 +10,7 @@
 
 def test_pytorch_estimator_creation():
     # Model
-    model = models.resnet18(pretrained=True)
+    model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
     model.fc = torch.nn.Linear(model.fc.in_features, 6)
 
     # Loss function