aapot
Add 400k train step and HF flax model
d51aeae