Ojas Patil comments

Results 5 comments of


                                            Ojas Patil

Add BatchNorm layers to CNN in MNIST tutorial for improved training stability

Can you merge this PR? @8bitmp3 @IvyZX @jburnim @cgarciae

Add BatchNorm layers to CNN in MNIST tutorial for improved training stability

> Can you merge this PR? @8bitmp3 @IvyZX @jburnim @cgarciae I’ve made the necessary changes to fix the GitHub Actions failure, everything checks out now.

Add BatchNorm layers to CNN in MNIST tutorial for improved training stability

thank you @cgarciae @8bitmp3 @IvyZX @jburnim for helping me with my first PR to FLAX

Pickling a trained model (NNX)

@cgarciae after using this [code ](https://github.com/google/flax/blob/main/examples/nnx_toy_examples/08_save_load_checkpoints.py)to save flax nnx model on colab `import orbax.checkpoint as orbax` `state = nnx.state(model)` `path = "/content/attention_model"` `checkpointer = orbax.PyTreeCheckpointer()` `checkpointer.save(f'{path}/state', state)` --------------------------------------------------------------------------- AttributeError Traceback...

Pickling a trained model (NNX)

Resolved the issue