update

modeling_yi.py

@@ -26,16 +26,13 @@ from .configuration_yi import YiConfig
 
 is_flash_attn_available = True
 try:
-    from flash_attn import flash_attn_func
-except Exception:
-    is_flash_attn_available = False
-
-if is_flash_attn_available:
-    from flash_attn import __version__
+    from flash_attn import flash_attn_func, __version__
 
     assert version.parse(__version__) >= version.parse(
         "2.3.0"
     ), "please update your flash_attn version (>= 2.3.0)"
+except ModuleNotFoundError:
+    is_flash_attn_available = False
 
 logger = logging.get_logger(__name__)