Update gemma_backbone_test.py

Better test messages
keras-team · Jun 20, 2024 · 2001a3d · 2001a3d
1 parent 12fc70c
commit 2001a3d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/keras_nlp/src/models/gemma/gemma_backbone_test.py b/keras_nlp/src/models/gemma/gemma_backbone_test.py
@@ -139,11 +139,11 @@ def test_distribution(self):
 
     def test_distribution_with_lora(self):
         if keras.backend.backend() != "jax":
-            return
+            self.skipTest("`ModelParallel` testing requires the Jax backend.")
         devices = keras.distribution.list_devices("CPU")
         if len(devices) == 1:
             # Need more than 1 device for distribution testing.
-            return
+            self.skipTest("`ModelParallel` testing requires multiple devices.")
         device_mesh = keras.distribution.DeviceMesh(
             shape=(1, len(devices)),
             axis_names=("batch", "model"),