Skip to content

Call start_prefetch immediately at the end of restoring#1211

Open
findmyway wants to merge 1 commit intogoogle:mainfrom
findmyway:patch-2
Open

Call start_prefetch immediately at the end of restoring#1211
findmyway wants to merge 1 commit intogoogle:mainfrom
findmyway:patch-2

Conversation

@findmyway
Copy link
Contributor

@findmyway findmyway commented Feb 1, 2026

The grain checkpoint handler is usually used together with other handlers (typically the array handler) in orbax under a composite handler. It is better to call start_prefetch immediately after setting state.


📚 Documentation preview 📚: https://google-grain--1211.org.readthedocs.build/

@iindyk iindyk self-assigned this Feb 2, 2026
copybara-service bot pushed a commit that referenced this pull request Feb 2, 2026
Ported from #1211

PiperOrigin-RevId: 864460741
copybara-service bot pushed a commit that referenced this pull request Feb 17, 2026
Ported from #1211

PiperOrigin-RevId: 864460741
copybara-service bot pushed a commit that referenced this pull request Feb 17, 2026
Ported from #1211

PiperOrigin-RevId: 864460741
copybara-service bot pushed a commit that referenced this pull request Feb 17, 2026
Ported from #1211

PiperOrigin-RevId: 864460741
copybara-service bot pushed a commit that referenced this pull request Feb 17, 2026
Ported from #1211

PiperOrigin-RevId: 871442961
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants