flax

59280a36 - Merge pull request #2457 from cgarciae:fix-checkpoint