added back skipping to test_pipeline_summary_writer

patrick-toulme · Dec 12, 2024 · 3d0ccab · 3d0ccab
1 parent 4a78272
commit 3d0ccab
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/axlearn/common/inference_test.py b/axlearn/common/inference_test.py
@@ -771,6 +771,9 @@ def test_pipeline_summary_writer(
         global_batch_size: int,
         data_partition: DataPartitionType,
     ):
+        supported, reason = is_supported(platform, mesh_shape, inference_dtype, global_batch_size, data_partition)
+        if not supported:
+            pytest.skip(reason=reason)
         del platform  # only used by is_supported_platform().
         local_run = jax.process_count() == 1
         mesh_axis_names = ("data", "model")