flax

322b305f - Merge pull request #3007 from andsteing:imagenet_gpu