https://github.com/google/jax
Raw File
Tip revision: d73319c8a25596af21dab9a9c12a468e1f6242d3 authored by Colin Gaffney on 10 June 2022, 22:09:01 UTC
Introduce AsyncCheckpointManager, which allows commit operations to take place in a background thread while training may continue in the main thread. This class inherits from both CheckpointManager, and reuses much of its functionality, and AsyncWriteManager in jax.experimental, which provides needed async utilities.
Tip revision: d73319c
.bazelversion
5.1.1
back to top