Allow mask-free usage

[mranzinger]

I tried to use your implementation, and ended up having to make a few tweaks in order to get the code to work again when I'm not supplying attn_mask, which should be a valid mode.

Even when I did try to supply attn_mask, I ran into some problems where using += was breaking broadcasting rules (implicit with the attn_mask.unsqueeze(0) call for batch), and also needed the

attn = attn.permute(0, 2, 1, 3).flatten(2)

to get the attn tensor into the expected format.

[usryokousha]

That looks about right! Sorry I should have included a unit test!