Update flax/linen/transforms.py

Commit

4 years ago

Update flax/linen/transforms.py Co-authored-by: 8bitmp3 <19637339+8bitmp3@users.noreply.github.com>