Update modeling_chatglm.py

modeling_chatglm.py

@@ -157,7 +157,7 @@ class RotaryEmbedding(nn.Module):
         )
 
 
-# @torch.jit.script
+@torch.jit.script
 def apply_rotary_pos_emb(x: torch.Tensor, rope_cache: torch.Tensor) -> torch.Tensor:
     # x: [sq, b, np, hn]
     sq, b, np, hn = x.size(0), x.size(1), x.size(2), x.size(3)
@@ -236,7 +236,7 @@ class CoreAttention(torch.nn.Module):
             # Raw attention scores
 
             # [b, np, sq, sk]
-            output_size = (query_layer.size(0), query_layer.size(2), query_layer.size(1), key_layer.size(0))
+            output_size = (query_layer.size(1), query_layer.size(2), query_layer.size(0), key_layer.size(0))
 
             # [sq, b, np, hn] -> [sq, b * np, hn]
             query_layer = query_layer.view(output_size[2], output_size[0] * output_size[1], -1)