flax

c16e5937 - Merge pull request #287 from mohitreddy1996:refactor-imagenet