Fix unit-test

openvinotoolkit · Sep 26, 2024 · 368a91c · 368a91c
1 parent de02ab5
commit 368a91c
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/unit/algo/classification/backbones/test_timm.py b/tests/unit/algo/classification/backbones/test_timm.py
@@ -11,11 +11,11 @@
 
 class TestOTXEfficientNetV2:
     def test_forward(self):
-        model = TimmBackbone(backbone="tf_efficientnetv2_s.in21k")
+        model = TimmBackbone(model_name="tf_efficientnetv2_s.in21k")
         assert model(torch.randn(1, 3, 244, 244))[0].shape == torch.Size([1, 1280, 8, 8])
 
     def test_get_config_optim(self):
-        model = TimmBackbone(backbone="tf_efficientnetv2_s.in21k")
+        model = TimmBackbone(model_name="tf_efficientnetv2_s.in21k")
         assert model.get_config_optim([0.01])[0]["lr"] == 0.01
         assert model.get_config_optim(0.01)[0]["lr"] == 0.01
 
@@ -24,5 +24,5 @@ def test_check_pretrained_weight_download(self):
         if target.exists():
             shutil.rmtree(target)
         assert not target.exists()
-        TimmBackbone(backbone="tf_efficientnetv2_s.in21k", pretrained=True)
+        TimmBackbone(model_name="tf_efficientnetv2_s.in21k", pretrained=True)
         assert target.exists()