flax

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

#281

Merged