flax
Added keep_rngs arg to MHA to optionally store rngs
#4749
Merged

Commits
  • Added keep_rngs arg to MHA to optionally store rngs
    vfdev-5 committed 1 year ago
Loading