flax

5a1787b3 - Remove loss_scaling flag from imagenet benchmark