flax

8753263b - Add --model_dir flag as the training code requires this field to be set.