423A35C7/LLaMA-Factory
1
0
Fork 0
mirror of https://github.com/hiyouga/LLaMA-Factory.git synced 2025-10-16 00:28:10 +08:00
Code Issues Packages Projects Releases Wiki Activity

Update test_attention.py

Browse Source 
Former-commit-id: c2cc7a0f152aa14fc03ae413f4a9dc06742a29d7
This commit is contained in:
hoshi-hiyouga 2024-06-24 21:35:34 +08:00 committed by GitHub
parent 16e950454e
commit 5675c47a01
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tests/model/model_utils/test_attention.py
View File

@ -29,7 +29,7 @@ INFER_ARGS = {
def test_attention(): def test_attention():
attention_available = ["off"] attention_available = ["disabled"]
if is_torch_sdpa_available(): if is_torch_sdpa_available():
attention_available.append("sdpa") attention_available.append("sdpa")
@ -37,7 +37,7 @@ def test_attention():
attention_available.append("fa2") attention_available.append("fa2")
llama_attention_classes = { llama_attention_classes = {
"off": "LlamaAttention", "disabled": "LlamaAttention",
"sdpa": "LlamaSdpaAttention", "sdpa": "LlamaSdpaAttention",
"fa2": "LlamaFlashAttention2", "fa2": "LlamaFlashAttention2",
} }