merge

Author: dirmeier

README.md

@@ -165,7 +165,7 @@ metrics = nnx.MultiMetric(
   accuracy=nnx.metrics.Accuracy(),
   loss=nnx.metrics.Average("loss"),
 )
-hook =  hook_fn(metrics, val_iter, hook_every_n_steps)
+hook = hook_fn(metrics, val_iter, hook_every_n_steps)
 ```
 
 This creates a hook function `hook` that after `eval_every_n_steps` steps iterates over the validation set

examples/mnist_classification/main.py

@@ -8,10 +8,11 @@
 from flax import nnx
 from jax import random as jr
 from jax.experimental import mesh_utils
-from model import CNN, train_step, val_step
 
 from blaxbird import get_default_checkpointer, train_fn
 
+from model import CNN, train_step, val_step
+
 
 def get_optimizer(model, lr=1e-4):
   tx = optax.adamw(lr)