Upload modeling_fegeo_llama.py

modeling_fegeo_llama.py

@@ -898,8 +898,11 @@ from .configuration_fegeo_llama import LlamaConfig
 
 
 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
+    try:
+        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
+    except:
+        flash_attn_func, flash_attn_varlen_func, index_first_axis, pad_input, unpad_input = None, None, None, None, None
 
 
 logger = logging.get_logger(__name__)