aapot
Add 550k train step and HF flax model
7f3c160