From 8137522328f4acaceeec4490952a4709bf0f316b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=97=AB=E9=B9=8F=E5=85=A8?= Date: Thu, 11 Sep 2025 19:38:23 +0800 Subject: [PATCH] support npu_fused_infer_attention_score.default for auto_dispatch --- torch_npu/npu/graphs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torch_npu/npu/graphs.py b/torch_npu/npu/graphs.py index 7e21ce5ed9..ef4d0d498b 100644 --- a/torch_npu/npu/graphs.py +++ b/torch_npu/npu/graphs.py @@ -135,7 +135,7 @@ class _GraphDispatchMode(torch.utils._python_dispatch.TorchDispatchMode): return _GraphDispatchRecord(event=event, handle=handle, kwargs=kwargs_ref, args=tuple(args_ref), op_cache_entry=func) def __torch_dispatch__(self, func, types, args=(), kwargs=None): - if func.__name__ == "npu_fused_infer_attention_score": + if func.__name__ in ["npu_fused_infer_attention_score", "npu_fused_infer_attention_score.default"]: func_out = torch_npu.npu_fused_infer_attention_score.out self.update_schema(str(func_out.__name__), str(func_out._schema)) stream = torch_npu.npu.current_stream() -- Gitee