From a3e7f60ae68a9a836fd878a9745a28b62ad33881 Mon Sep 17 00:00:00 2001 From: li-lizhe <82501526+li-lizhe@users.noreply.github.com> Date: Tue, 3 Nov 2026 20:54:27 +0800 Subject: [PATCH] fix(longlora): extend device type check to include Ascend NPU (#10829) --- src/llamafactory/model/model_utils/longlora.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/llamafactory/model/model_utils/longlora.py b/src/llamafactory/model/model_utils/longlora.py index f7c36ee0c..07bf8f5df 100644 --- a/src/llamafactory/model/model_utils/longlora.py +++ b/src/llamafactory/model/model_utils/longlora.py @@ -317,7 +317,7 @@ def llama_sdpa_attention_forward( if attention_mask is not None: 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() key_states = key_states.contiguous() value_states = value_states.contiguous()