Update modeling_gritlm7b.py

modeling_gritlm7b.py

@@ -46,12 +46,13 @@ from ...utils import (
 from .configuration_mistral import MistralConfig
 
 
-if is_flash_attn_2_available():
+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
 
     _flash_supports_window_size = "window_size" in list(inspect.signature(flash_attn_func).parameters)
-
+except:
+    pass
 
 logger = logging.get_logger(__name__)