Skip the train module only.

tests/models/autoencoder_linear/test_autoencoder_linear.py

@@ -84,7 +84,8 @@ def _load_inputs(self):
     ["train", "eval"],
 )
 def test_autoencoder_linear(record_property, mode):
-    pytest.skip("module has train variant.")
+    if mode == "train":
+        pytest.skip()
     model_name = "Autoencoder (linear)"
     record_property("model_name", model_name)
     record_property("mode", mode)

tests/models/beit/test_beit_image_classification.py

@@ -43,7 +43,8 @@ def get_results_train(self, model, inputs, outputs):
     ["microsoft/beit-base-patch16-224", "microsoft/beit-large-patch16-224"],
 )
 def test_beit_image_classification(record_property, model_name, mode):
-    pytest.skip("module has train variant.")
+    if mode == "train":
+        pytest.skip()
     record_property("model_name", model_name)
     record_property("mode", mode)
 

tests/models/clip/test_clip.py

@@ -59,7 +59,8 @@ def get_results_train(self, model, inputs, outputs):
     ],
 )
 def test_clip(record_property, mode):
-    pytest.skip("module has train variant.")
+    if mode == "train":
+        pytest.skip()
     model_name = "CLIP"
     record_property("model_name", model_name)
     record_property("mode", mode)

tests/models/deit/test_deit.py

@@ -47,7 +47,8 @@ def get_results_train(self, model, inputs, outputs):
 )
 @pytest.mark.parametrize("model_name", ["facebook/deit-base-patch16-224"])
 def test_deit(record_property, model_name, mode):
-    pytest.skip("module has train variant.")
+    if mode == "train":
+        pytest.skip()
     record_property("model_name", model_name)
     record_property("mode", mode)
 

tests/models/hardnet/test_hardnet.py

@@ -50,7 +50,8 @@ def _load_inputs(self):
     ["train", "eval"],
 )
 def test_hardnet(record_property, mode):
-    pytest.skip("module has train variant.")
+    if mode == "train":
+        pytest.skip()
     model_name = "HardNet"
     record_property("model_name", model_name)
     record_property("mode", mode)

tests/models/mlpmixer/test_mlpmixer.py

@@ -34,7 +34,8 @@ def _load_inputs(self):
     ["train", "eval"],
 )
 def test_mlpmixer(record_property, mode):
-    pytest.skip("module has train variant.")
+    if mode == "train":
+        pytest.skip()
     model_name = "MLPMixer"
     record_property("model_name", model_name)
     record_property("mode", mode)

tests/models/mnist/test_mnist.py

@@ -60,7 +60,8 @@ def _load_inputs(self):
     ["train", "eval"],
 )
 def test_mnist_train(record_property, mode):
-    pytest.skip("module has train variant.")
+    if mode == "train":
+        pytest.skip()
     model_name = "Mnist"
     record_property("model_name", model_name)
     record_property("mode", mode)

tests/models/openpose/test_openpose_v2.py

@@ -50,7 +50,8 @@ def _load_inputs(self):
     ["train", "eval"],
 )
 def test_openpose_v2(record_property, mode):
-    pytest.skip("module has train variant.")
+    if mode == "train":
+        pytest.skip()
     model_name = "OpenPose V2"
     record_property("model_name", model_name)
     record_property("mode", mode)

tests/models/resnet/test_resnet.py

@@ -24,7 +24,8 @@ def _load_inputs(self):
     ["train", "eval"],
 )
 def test_resnet(record_property, mode):
-    pytest.skip("module has train variant.")
+    if mode == "train":
+        pytest.skip()
     model_name = "ResNet18"
     record_property("model_name", model_name)
     record_property("mode", mode)

tests/models/resnet50/test_resnet50.py

@@ -37,7 +37,8 @@ def _load_inputs(self):
     ["train", "eval"],
 )
 def test_resnet(record_property, mode):
-    pytest.skip("module has train variant.")
+    if mode == "train":
+        pytest.skip()
     model_name = "ResNet50"
     record_property("model_name", model_name)
     record_property("mode", mode)

tests/models/segformer/test_segformer.py

@@ -45,7 +45,8 @@ def get_results_train(self, model, inputs, outputs):
     ["train", "eval"],
 )
 def test_segformer(record_property, mode):
-    pytest.skip("module has train variant.")
+    if mode == "train":
+        pytest.skip()
     model_name = "SegFormer"
     record_property("model_name", model_name)
     record_property("mode", mode)

tests/models/speecht5_tts/test_speecht5_tts.py

@@ -48,7 +48,7 @@ def set_model_eval(self, model):
     ["eval"],
 )
 def test_speecht5_tts(record_property, mode):
-    pytest.skip("crashes in lowering to stable hlo.")
+    pytest.skip()  # crashes in lowering to stable hlo
     model_name = "speecht5-tts"
     record_property("model_name", model_name)
     record_property("mode", mode)

tests/models/unet/test_unet.py

@@ -46,7 +46,8 @@ def _load_inputs(self):
     ["train", "eval"],
 )
 def test_unet(record_property, mode):
-    pytest.skip("module has train variant.")
+    if mode == "train":
+        pytest.skip()
     model_name = "U-Net"
     record_property("model_name", model_name)
     record_property("mode", mode)

tests/models/unet_brain/test_unet_brain.py

@@ -56,7 +56,8 @@ def _load_inputs(self):
     ["train", "eval"],
 )
 def test_unet_brain(record_property, mode):
-    pytest.skip("module has train variant.")
+    if mode == "train":
+        pytest.skip()
     model_name = "Unet-brain"
     record_property("model_name", model_name)
     record_property("mode", mode)

tests/models/unet_carvana/test_unet_carvana.py

@@ -32,7 +32,8 @@ def _load_inputs(self):
     ["train", "eval"],
 )
 def test_unet_carvana(record_property, mode):
-    pytest.skip("module has train variant.")
+    if mode == "train":
+        pytest.skip()
     model_name = "Unet-carvana"
     record_property("model_name", model_name)
     record_property("mode", mode)