orby-yanan's picture
Allow returning attention weights
1662aa8