PR #6572 Add optimizer priming for dist chkpt

Add optimizer priming for dist chkpt #6572

jonb377 merged 4 commits into master from jonbolin/prime

jonb377 requested a review from

yeounoh 2 years ago

jonb377 requested a review from

alanwaketan 2 years ago

jonb377 force pushed from 82dfbe2a to 1e40c744 2 years ago

alanwaketan approved these changes on 2024-02-23

Add optimizer priming for dist chkpt

4e9aae0d

Only run test on TPUs

ae713fcf

Remove adagrad test

b7044e84

jonb377 force pushed from 98c51e9b to b7044e84 2 years ago

Wait for device ops before returning

3d142329

jonb377 merged 8fd051f2 into master 2 years ago

jonb377 deleted the jonbolin/prime branch 2 years ago

Reviewers

alanwaketan

yeounoh

Assignees

No one assigned

Labels

None yet

Milestone

No milestone