Update bert_layers.py

bert_layers.py

@@ -485,7 +485,8 @@ class BertEncoder(nn.Module):
                                                               attn_mask=attention_mask,
                                                               bias=alibi_attn_mask)
             all_attention_weights.append(attention_weights)  # Store attention weights
-            print(all_attention_weights)
+            print(f'here is the matrix of attentions inside encoder: \n {all_attention_weights}')
+            print(f'and this is the shape inside encoder: \n {all_attention_weights.shape}')
     
         if not output_all_encoded_layers:
             all_encoder_layers.append(hidden_states)
@@ -613,11 +614,13 @@ class BertModel(BertPreTrainedModel):
             first_col_mask[:, 0] = True
             subset_mask = masked_tokens_mask | first_col_mask
 
-        encoder_outputs = self.encoder(
+        encoder_outputs, all_attentions = self.encoder(
             embedding_output,
             attention_mask,
             output_all_encoded_layers=output_all_encoded_layers,
             subset_mask=subset_mask)
+        print(f'here is the matrix of attentions in BERT: \n {all_attention_weights}')
+        print(f'and this is the shape in BERT: \n {all_attention_weights.shape}')
 
         if masked_tokens_mask is None:
             sequence_output = encoder_outputs[-1]