Update modeling_gritlm7b.py

modeling_gritlm7b.py

@@ -46,11 +46,13 @@ from transformers.utils import (
 from transformers import MistralConfig
 
 
+# transformers has a bug where it will try to import everything from a custom model file unless there's try/except
 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)
+    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
+    
+        _flash_supports_window_size = "window_size" in list(inspect.signature(flash_attn_func).parameters)
 except:
     pass