flax

aa299206 - Merge pull request #717 from andsteing:imagenet_bugfix