ayo/ipex-llm
1
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

[LLM] Fix dtype mismatch in Baichuan2-13b (#9834)

Browse source
This commit is contained in:
Xiangyu Tian 2024-01-04 15:34:42 +08:00 committed by GitHub
parent 8504a2bbca
commit 38c05be1c0
1 changed files with 1 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
python/llm/src/bigdl/llm/transformers/models/baichuan2.py
View file

@ -287,7 +287,7 @@ def baichuan_attention_forward_13b(
)
attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1)
attn_output = torch.matmul(attn_weights, value_states)
attn_output = torch.matmul(attn_weights.to(dtype=value_states.dtype), value_states)
attn_output = attn_output.transpose(1, 2)
attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)