Update bert_layers.py

bert_layers.py

@@ -491,6 +491,7 @@ class BertEncoder(nn.Module):
         if not output_all_encoded_layers:
             all_encoder_layers.append(hidden_states)
         return all_encoder_layers, all_attention_weights  # Return both hidden states and attention weights
+        return all_encoder_layers  # Return both hidden states and attention weights
 
 
 
@@ -614,13 +615,13 @@ class BertModel(BertPreTrainedModel):
             first_col_mask[:, 0] = True
             subset_mask = masked_tokens_mask | first_col_mask
 
-        encoder_outputs, all_attentions = self.encoder(
+        encoder_outputs = 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}')
+        # 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]