Solved by tom in post #2 
              
              
                You cannot use indexed assignment here, as the backward of the previous operation wants to use the calculation result you are overwriting. 
The typical solution is to use torch.where(mask, attn_weights, neginf) where neginf is torch.tensor(float('-inf'), device=device) (I would recommend precomputin…