Fix test_gpu_distributed_initialize() to get a more proper GPUs number (#503)

Author: Arech8

tests/multiprocess_gpu_test.py

@@ -50,7 +50,7 @@ def test_gpu_distributed_initialize(self):
 
     port = portpicker.pick_unused_port()
     # ROCm will fail if HIP_VISIBLE_DEVICES is set to a single non existent device
-    num_gpus = 4 if not jtu.is_device_rocm() else min(4, len(jax.devices()))
+    num_gpus = 4 if not jtu.is_device_rocm() else min(4, jax.local_device_count())
     num_gpus_per_task = 1
     num_tasks = num_gpus // num_gpus_per_task