Upload modeling_esm_plusplus.py with huggingface_hub

modeling_esm_plusplus.py

@@ -316,12 +316,11 @@ class MultiHeadAttention(nn.Module):
         query_BHLD, key_BHLD, value_BHLD = map(self.reshaper, (query_BLD, key_BLD, value_BLD))
 
         if output_attentions: # Manual attention computation
-            b, L, d = x.shape
+            b, h, l, d = query_BHLD.shape
             scale = 1 / math.sqrt(d)
-            attn_bias = torch.zeros(b, 1, L, L, dtype=query_BLD.dtype, device=query_BLD.device)
+            attn_bias = torch.zeros(b, h, l, l, dtype=query_BLD.dtype, device=query_BLD.device)
             if attention_mask is not None:
                 attn_bias.masked_fill_(attention_mask.logical_not(), float('-inf'))
-
             attn_weights = torch.matmul(query_BHLD, key_BHLD.transpose(-2, -1)) * scale
             attn_weights += attn_bias
             attn_weights = F.softmax(attn_weights, dim=-1)