Update bert_layers.py

bert_layers.py

@@ -199,9 +199,9 @@ class BertUnpadSelfAttention(nn.Module):
 
         # print(f'PROBLEM HERE: UNDERSTAND IT!!')
         rearranged_attention = rearrange(attention, 'nnz h d -> nnz (h d)')
-        try:
+        # try:
             # print(f'REARRANGED ATTENTION: {rearranged_attention.shape}')
-        except:
+        # except:
             # print(f'REARRANGED ATTENTION: {rearranged_attention[0].shape}')
         return rearrange(attention, 'nnz h d -> nnz (h d)'), attention_probs