bad commit for jax devices

youliangtan · Nov 1, 2023 · b6d7460 · b6d7460
1 parent fccc37d
commit b6d7460
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/async_learner_actor.py b/examples/async_learner_actor.py
@@ -222,7 +222,7 @@ def stats_callback(type: str, payload: dict) -> dict:
 
 
 def main(_):
-    devices = jax.local_devices()[:1]
+    devices = jax.local_devices()
     num_devices = len(devices)
     sharding = jax.sharding.PositionalSharding(devices)
     assert FLAGS.batch_size % num_devices == 0