Update bert_layers.py

bert_layers.py

@@ -248,6 +248,12 @@ class BertUnpadAttention(nn.Module):
         """
         self_output = self.self(input_tensor, cu_seqlens, max_s, indices,
                                 attn_mask, bias)
+
+        try:
+            print(f'IMPORTANT: {self_output.shape}')
+        except:
+            print(f'IMPORTANT2: {self_output[0].shape}')
+            
         if subset_idx is not None:
             return self.output(index_first_axis(self_output, subset_idx),
                                index_first_axis(input_tensor, subset_idx))