Fix local restore by re-mapping device ids directly instead of inferring them from how process indexes changed across restarts with some false assumptions.

google / orbax

Orbax provides common checkpointing and persistence utilities for JAX users

https://orbax.readthedocs.io/

Apache License 2.0

305 stars 36 forks source link

Fix local restore by re-mapping device ids directly instead of inferring them from how process indexes changed across restarts with some false assumptions. #1392

Open copybara-service[bot] opened 21 hours ago

copybara-service[bot] commented 21 hours ago

Fix local restore by re-mapping device ids directly instead of inferring them from how process indexes changed across restarts with some false assumptions.