flax

8b871506 - Fixes for batch norm docs

Fixes for batch norm docs Changes: - Fixed the docstring in `eval_step` - Added `.mean()` to losses - Fixed reference to `params` in `eval_step`