flax

bfe304cd - Merge pull request #563 from jpuigcerver:checkpoint-restore-target-none