p12475398
/
LLaMA-Factory-Mirror
forked from p04798526/LLaMA-Factory-Mirror
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update test_attention.py

This commit is contained in:
hoshi-hiyouga 2024-06-24 21:35:34 +08:00 committed by GitHub
parent 3ed063f281
commit a9b3d91952
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
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():
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",
}