Update jax/_src/xla_bridge.py

Co-authored-by: Yash Katariya <yashkatariya@google.com>
jax-ml · Mar 13, 2024 · 2d76642 · 2d76642
1 parent ac2c522
commit 2d76642
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/jax/_src/xla_bridge.py b/jax/_src/xla_bridge.py
@@ -336,8 +336,8 @@ def make_gpu_client(
     _check_cuda_versions()
     # TODO(micky774): remove this check when minimum jaxlib is v0.4.26
     if version.__version_info__ >= (0, 4, 26):
-      devices_to_check = allowed_devices if allowed_devices else range(
-        cuda_versions.cuda_device_count())
+      devices_to_check = (allowed_devices if allowed_devices
+                                          else range(cuda_versions.cuda_device_count()))
       _check_cuda_compute_capability(devices_to_check)
 
   return xla_client.make_gpu_client(