use packaging.version

modeling_yi.py

@@ -4,6 +4,7 @@ from typing import List, Optional, Tuple, Union
 
 import torch.utils.checkpoint
 from einops import repeat
+from packaging import version
 from torch import nn
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
 from transformers.activations import ACT2FN
@@ -32,7 +33,9 @@ except Exception:
 if is_flash_attn_available:
     from flash_attn import __version__
 
-    assert __version__ >= "2.3.0", "please update your flash_attn version (>= 2.3.0)"
+    assert version.parse(__version__) >= version.parse(
+        "2.3.0"
+    ), "please update your flash_attn version (>= 2.3.0)"
 
 logger = logging.get_logger(__name__)