Update modeling_intern_vit.py (#4)

- Update modeling_intern_vit.py (188f686e8863f48f045de9508e6035f37809e1fd)


Co-authored-by: Roy Hvaara <[email protected]>

modeling_intern_vit.py

@@ -15,17 +15,18 @@ from transformers.activations import ACT2FN
 from transformers.modeling_outputs import (BaseModelOutput,
                                            BaseModelOutputWithPooling)
 from transformers.modeling_utils import PreTrainedModel
+from transformers.utils.import_utils import is_flash_attn_greater_or_equal
 from transformers.utils import logging
 
 from .configuration_intern_vit import InternVisionConfig
 
 try:
-    try:  # v1
-        from flash_attn.flash_attn_interface import \
-            flash_attn_unpadded_qkvpacked_func
-    except:  # v2
+    if is_flash_attn_greater_or_equal("2.0.0"):
         from flash_attn.flash_attn_interface import \
             flash_attn_varlen_qkvpacked_func as flash_attn_unpadded_qkvpacked_func
+    else:
+        from flash_attn.flash_attn_interface import \
+            flash_attn_unpadded_qkvpacked_func
 
     from flash_attn.bert_padding import pad_input, unpad_input