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

Update test_attention.py

Browse Source
This commit is contained in:
hoshi-hiyouga
2024-06-24 21:35:34 +08:00
committed by GitHub
parent 3ed063f281
commit a9b3d91952
1 changed files with 2 additions and 2 deletions
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():
attention_available = ["off"]
attention_available = ["disabled"]
if is_torch_sdpa_available():
attention_available.append("sdpa")
@@ -37,7 +37,7 @@ def test_attention():
attention_available.append("fa2")
llama_attention_classes = {
"off": "LlamaAttention",
"disabled": "LlamaAttention",
"sdpa": "LlamaSdpaAttention",
"fa2": "LlamaFlashAttention2",
}