Fix tensor shape error (#7)

- Fix tensor shape error (3d222b84ea089f5c9463cb7e7327b81a221fa4ec)


Co-authored-by: Yaowei Zheng <hiyouga@users.noreply.huggingface.co>

modeling_chatglm.py

@@ -253,15 +253,12 @@ class CoreAttention(torch.nn.Module):
             # This is actually dropping out entire tokens to attend to, which might
             # seem a bit unusual, but is taken from the original Transformer paper.
             attention_probs = self.attention_dropout(attention_probs)
-            # =========================
-            # Context layer. [sq, b, hp]
-            # =========================
-
-            # value_layer -> context layer.
-            # [sk, b, np, hn] --> [b, np, sq, hn]
 
+            # query layer shape: [b * np, sq, hn]
+            # value layer shape: [b, np, sk, hn]
+            # attention shape: [b, np, sq, sk]
             # context layer shape: [b, np, sq, hn]
-            output_size = (value_layer.size(1), value_layer.size(2), query_layer.size(0), value_layer.size(3))
+            output_size = (value_layer.size(0), value_layer.size(1), query_layer.size(1), value_layer.size(3))
             # change view [b * np, sk, hn]
             value_layer = value_layer.view(output_size[0] * output_size[1], value_layer.size(2), -1)
             # change view [b * np, sq, sk]