Disable multiprocess_gpu_test that fails on ROCm.

google · Sep 28, 2022 · f7bafb3 · f7bafb3
1 parent 96abd9a
commit f7bafb3
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tests/multiprocess_gpu_test.py b/tests/multiprocess_gpu_test.py
@@ -129,6 +129,7 @@ def test_gpu_distributed_initialize(self):
 
   @unittest.skipIf(xla_extension_version < 91,
                    "Test requires jaxlib 0.3.17 or newer")
+  @jtu.skip_on_devices("rocm")
   def test_distributed_jax_cuda_visible_devices(self):
     """Test jax_cuda_visible_devices works in distributed settings."""
     if jax.devices()[0].platform != 'gpu':