OrionStarAI
/

Orion-14B-Base

Text Generation

Model card Files Files and versions

Update modeling_orion.py

#5

by goldbach7 - opened Jul 10, 2024

base: refs/heads/main

←

from: refs/pr/5

Discussion Files changed

Files changed (1) hide show

modeling_orion.py +2 -2

modeling_orion.py CHANGED Viewed

@@ -25,12 +25,12 @@ from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
 from transformers.utils import (
     add_start_docstrings,
     add_start_docstrings_to_model_forward,
-    is_flash_attn_available,
     logging,
     replace_return_docstrings,
 )
-if is_flash_attn_available():
     from flash_attn import flash_attn_func, flash_attn_varlen_func
     from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa

 from transformers.utils import (
     add_start_docstrings,
     add_start_docstrings_to_model_forward,
+    is_flash_attn_2_available,
     logging,
     replace_return_docstrings,
 )
+if is_flash_attn_2_available():
     from flash_attn import flash_attn_func, flash_attn_varlen_func
     from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa