search

apple / axlearn

An Extensible Deep Learning Library
Apache License 2.0
1.88k stars 269 forks source link

[Bug Fix] Fix a nit issue for Jax rollback #856

Closed kelvin-zou closed 1 day ago

kelvin-zou commented 3 days ago

Already verified by internal TPU job.