[torchax][distributed] Use local device for default jax device (#8667)

Co-authored-by: Siyuan Liu <lsiyuan@google.coim>
pytorch · Feb 3, 2025 · 3578940 · 3578940
1 parent 9ae017e
commit 3578940
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchax/torchax/tensor.py b/torchax/torchax/tensor.py
@@ -304,7 +304,7 @@ def get_as_jax_device(self, device: Any):
 
       if device == 'cpu':
         return jax.devices('cpu')[0]
-      return jax.devices()[0]
+      return jax.local_devices()[0]
 
 
     def load_ops(self):