Tests

Author: qubvel

tests/models/base.py

@@ -33,6 +33,8 @@ class BaseModelTester(unittest.TestCase):
     default_height = 64
     default_width = 64
 
+    compile_dynamic = True
+
     @property
     def model_type(self):
         if self.test_model_type is None:
@@ -232,16 +234,16 @@ def test_compile(self):
         model = model.eval().to(default_device)
 
         if not model._is_torch_compilable:
-            with self.assertRaises(RuntimeError):
+            with self.assertRaises((RuntimeError)):
                 torch.compiler.reset()
                 compiled_model = torch.compile(
-                    model, fullgraph=True, dynamic=True, backend="eager"
+                    model, fullgraph=True, dynamic=self.compile_dynamic, backend="eager"
                 )
             return
 
         torch.compiler.reset()
         compiled_model = torch.compile(
-            model, fullgraph=True, dynamic=True, backend="eager"
+            model, fullgraph=True, dynamic=self.compile_dynamic, backend="eager"
         )
         with torch.inference_mode():
             compiled_model(sample)

tests/models/test_dpt.py

@@ -20,15 +20,20 @@ class TestDPTModel(base.BaseModelTester):
     # should be overriden
     test_model_type = "dpt"
 
+    compile_dynamic = False
+
     @property
     def hub_checkpoint(self):
-        return "smp-hub/dpt-large-ade20k"
+        return "smp-test-models/dpt-tu-test_vit"
 
     @slow_test
     @requires_torch_greater_or_equal("2.0.1")
     @pytest.mark.logits_match
-    def test_preserve_forward_output(self):
-        model = smp.from_pretrained(self.hub_checkpoint).eval().to(default_device)
+    def test_load_pretrained(self):
+        hub_checkpoint = "smp-hub/dpt-large-ade20k"
+
+        model = smp.from_pretrained(hub_checkpoint)
+        model = model.eval().to(default_device)
 
         input_tensor = torch.ones((1, 3, 384, 384))
         input_tensor = input_tensor.to(default_device)