Fix for non-numerical version number in flash-attn

turboderp · Sep 19, 2023 · ec5164b · ec5164b
1 parent b9cbe23
commit ec5164b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/exllamav2/attn.py b/exllamav2/attn.py
@@ -16,7 +16,7 @@
 has_flash_attn = False
 try:
     import flash_attn
-    flash_attn_ver = [int(t) for t in flash_attn.__version__.split(".")]
+    flash_attn_ver = [int(t) for t in flash_attn.__version__.split(".") if t.isdigit()]
     if flash_attn_ver >= [2, 2, 1]:
         from flash_attn import flash_attn_func
         has_flash_attn = True