Update bert_layers.py

bert_layers.py

@@ -195,6 +195,11 @@ class BertUnpadSelfAttention(nn.Module):
         # attn_mask is 1 for attend and 0 for don't
         attention = unpad_input_only(attention, torch.squeeze(attn_mask) == 1)
         print(f'BUSA unpadded final attention shape: {attention_probs.shape}')
+        rearranged_attention = rearrange(attention, 'nnz h d -> nnz (h d)')
+        try:
+            print(f'REARRANGED ATTENTION: {rearranged_attention.shape}')
+        except:
+            print(f'REARRANGED ATTENTION: {rearranged_attention[0].shape}')
         return rearrange(attention, 'nnz h d -> nnz (h d)')