multihead attention version with loop by batch dimension to reduce memory usage#10

krunt wants to merge 4 commits intolearning-at-home:mainlearning-at-home/lean_transformer:mainfrom

krunt:mainkrunt/lean_transformer:mainCopy head branch name to clipboard