[sciences44]

Nice work. Regarding your first insight, that's a sneaky one ^^. I experiment a similar issue with fine-tuning a small model with MLX. It was very hard to find out.

[tucan9389]

Thank you. Agreed, we should be more careful when doing this. Was it the JAX training codebase?