Solved by tumble-weed in post #2 
              
              
                you can add singleton dimensions using None in attention_mask 
attention_logits = attention_logits.masked_fill(attention_mask[:,None,:,None],-1e7)