test: Added dedicated unittest

frgfm · Aug 1, 2022 · 4a0f35f · 4a0f35f
1 parent 8a92588
commit 4a0f35f
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/tests/test_methods_utils.py b/tests/test_methods_utils.py
@@ -1,4 +1,4 @@
-from torchvision.models import resnet18
+from torchvision.models import resnet18, mobilenet_v3_large
 
 from torchcam.methods import _utils
 
@@ -8,6 +8,10 @@ def test_locate_candidate_layer(mock_img_model):
     mod = resnet18().eval()
     assert _utils.locate_candidate_layer(mod) == "layer4"
 
+    # Mobilenet V3 Large
+    mod = mobilenet_v3_large().eval()
+    assert _utils.locate_candidate_layer(mod) == "features"
+
     # Custom model
     mod = mock_img_model.train()