mirror of
https://github.com/hiyouga/LLaMA-Factory.git
synced 2026-11-06 15:10:01 +08:00
fix(longlora): extend device type check to include Ascend NPU (#10829)
This commit is contained in:
@@ -317,7 +317,7 @@ def llama_sdpa_attention_forward(
|
|||||||
if attention_mask is not None:
|
if attention_mask is not None:
|
||||||
causal_mask = causal_mask[:, :, :, : key_states.shape[-2]]
|
causal_mask = causal_mask[:, :, :, : key_states.shape[-2]]
|
||||||
|
|
||||||
if query_states.device.type == "cuda" and causal_mask is not None: # avoid pytorch bug
|
if query_states.device.type in ("cuda", "npu") and causal_mask is not None: # avoid pytorch bug
|
||||||
query_states = query_states.contiguous()
|
query_states = query_states.contiguous()
|
||||||
key_states = key_states.contiguous()
|
key_states = key_states.contiguous()
|
||||||
value_states = value_states.contiguous()
|
value_states = value_states.contiguous()
|
||||||
|
|||||||
Reference in New Issue
Block a user