Update docs/guides/flax_basics.md

Commit

3 years ago

Update docs/guides/flax_basics.md Co-authored-by: 8bitmp3 <19637339+8bitmp3@users.noreply.github.com>