flax

95409ece - Merge pull request #1200 from hmph:checkpoint-fix