Update bert_layers.py

bert_layers.py

@@ -950,14 +950,18 @@ class BertForSequenceClassification(BertPreTrainedModel):
             return ((loss,) + output) if loss is not None else output
             
         # print(outputs.attentions)
-        # try:
-        #     print(f'not stacked final attention SHAPE: {outputs[2][0].shape}')
-        # except:
-        #     print(f'not stacked final attention LEN: {len(outputs[2])}')
-        # try:
-        #     print(f'STACKED final attention SHAPE: {(torch.stack(outputs[2], dim=0)).shape}')
-        # except:
-        #     print(f'STACKED final attention LEN: {len(torch.stack(outputs[2], dim=0))}')
+        try:
+            print(f'not stacked final attention SHAPE: {outputs[2][0].shape}')
+        except:
+            print(f'not stacked final attention LEN: {len(outputs[2])}')
+            
+        try:
+            print(f'STACKED final attention SHAPE: {(outputs.attentions).shape}')
+        except:
+            try:
+                print(f'STACKED final attention LEN: {(outputs.attentions)[0].shape}')
+            except:
+                print(f'STACKED final attention LEN 2: {len(outputs.attentions)}")
             
         return SequenceClassifierOutput(
             loss=loss,