flax
changed `return_weights` to `sow_weights` for attention layer
#3550
Merged

Commits
  • changed to for attention layer
    chiamp committed 2 years ago
Loading