feat: Removes torchvision warnings

frgfm · Oct 19, 2023 · 72d9e33 · 72d9e33
1 parent 88eb30f
commit 72d9e33
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git a/tests/test_methods_activation.py b/tests/test_methods_activation.py
@@ -6,7 +6,7 @@
 
 
 def test_base_cam_constructor(mock_img_model):
- model = mobilenet_v2(pretrained=False).eval()
+ model = mobilenet_v2(weights=None).eval()
  for p in model.parameters():
  p.requires_grad_(False)
  # Check that multiple target layers is disabled for base CAM
@@ -39,7 +39,7 @@ def _verify_cam(activation_map, output_size):
  ],
 )
 def test_img_cams(cam_name, target_layer, fc_layer, num_samples, output_size, batch_size, mock_img_tensor):
- model = mobilenet_v2(pretrained=False).eval()
+ model = mobilenet_v2(weights=None).eval()
  for p in model.parameters():
  p.requires_grad_(False)
  kwargs = {}

diff --git a/tests/test_methods_gradient.py b/tests/test_methods_gradient.py
@@ -26,7 +26,7 @@ def _verify_cam(activation_map, output_size):
  ],
 )
 def test_img_cams(cam_name, target_layer, output_size, batch_size, mock_img_tensor):
- model = mobilenet_v2(pretrained=False).eval()
+ model = mobilenet_v2(weights=None).eval()
  for p in model.parameters():
  p.requires_grad_(False)
 
@@ -79,7 +79,7 @@ def test_video_cams(cam_name, target_layer, output_size, mock_video_model, mock_
 
 
 def test_smoothgradcampp_repr():
- model = mobilenet_v2(pretrained=False).eval()
+ model = mobilenet_v2(weights=None).eval()
 
  # Hook the corresponding layer in the model
  with gradient.SmoothGradCAMpp(model, "features.18.0") as extractor:

diff --git a/tests/test_metrics.py b/tests/test_metrics.py
@@ -8,7 +8,7 @@
 
 
 def test_classification_metric():
- model = mobilenet_v3_small(pretrained=False)
+ model = mobilenet_v3_small(weights=None)
  with LayerCAM(model, "features.12") as extractor:
  metric = metrics.ClassificationMetric(extractor, partial(torch.softmax, dim=-1))