LLM: fix llama 4.36 forward(#10047)

2024-01-31 10:31:10 +08:00 · 2024-01-31 10:31:10 +08:00 · 3685622f29
commit 3685622f29
parent 53a5140eff
1 changed files with 1 additions and 1 deletions
--- a/python/llm/src/bigdl/llm/transformers/models/llama.py
+++ b/python/llm/src/bigdl/llm/transformers/models/llama.py
@ -638,7 +638,7 @@ def llama_attention_forward_4_36(
            "Please make sure use `attention_mask` instead.`"
        )

-    bsz, q_len, _ = hidden_states.size()
+    bsz, q_len, hidden_size = hidden_states.size()
    device = hidden_states.device
    # for flash attention
    original_dtype = hidden_states.dtype