Update bert_layers.py

bert_layers.py

@@ -903,6 +903,7 @@ class BertForSequenceClassification(BertPreTrainedModel):
         
         # JAANDOUI:
         all_attention_weights = outputs[2]
+        
         # print(f'last: {all_attention_weights}')
 
         pooled_output = self.dropout(pooled_output)
@@ -947,6 +948,6 @@ class BertForSequenceClassification(BertPreTrainedModel):
             logits=logits,
             hidden_states=outputs[0],
             #JAANDOUI: returning all_attention_weights here
-            attentions=outputs[2],
+            attentions=torch.stack(outputs[2], dim=0),
         )