Update bert_layers.py

bert_layers.py

@@ -196,6 +196,8 @@ class BertUnpadSelfAttention(nn.Module):
         attention = unpad_input_only(attention, torch.squeeze(attn_mask) == 1)
         print(f'BUSA unpadded final attention shape: {attention_probs.shape}')
         print(f'ATTENTION: {attention.shape}')
+
+        print(f'PROBLEM HERE: UNDERSTAND IT!!')
         rearranged_attention = rearrange(attention, 'nnz h d -> nnz (h d)')
         try:
             print(f'REARRANGED ATTENTION: {rearranged_attention.shape}')