flax

Update imagenet_fake_data_benchmark to include warmup phase

#1466

Merged