Update modeling_orion.py

modeling_orion.py

@@ -25,7 +25,7 @@ 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, 
+    is_flash_attn_2_available, 
     logging,
     replace_return_docstrings,
 )
@@ -33,7 +33,7 @@ from .generation_utils import build_chat_input, TextIterStreamer
 from transformers.generation.utils import GenerationConfig
 from threading import Thread
 
-if is_flash_attn_available():
+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