flax

f14874d3 - Benchmarks should run main func instead of executing the training loop directly.