def METHOD_NAME(self, attn_weights, tgt_len, src_len, bsz):
    sparse_mask = self.buffered_sparse_mask(attn_weights, tgt_len, src_len)
    sparse_mask = sparse_mask.unsqueeze(0).expand(
        bsz * self.num_heads, tgt_len, src_len
    )
    attn_weights += sparse_mask