Fixes for batch norm docs

[jkarwowski]

Changes:

• Fixed reference to params in eval_step
• Fixed the docstring in eval_step
• Added .mean() to losses

The documentation for batch norm is written in such a way, that directly copying from it results in an unpleasant bug where the params used in eval_step refer to the initial parameters of the network. This causes metrics to show that the model is not generalising to the test set at all. This change hopefully saves someone else learning flax a similar scratch on the head.

Checklist

• This PR fixes a minor issue (e.g.: typo or small bug) or improves the docs (you can dismiss the other
  checks if that's the case).

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[cgarciae]

Thanks @jkarwowski for doing this!