upload flax model

allow flax

Retrained model with ReLU

Change activation to use ReLU

Fix config

Initial commit

initial commit