File size: 907 Bytes
f545704 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 | import importlib.util
import torch
def is_flash_attn_available():
return importlib.util.find_spec("flash_attn") is not None
def is_flash_attn_3_available():
return importlib.util.find_spec("flash_attn_interface") is not None
def is_torch_version(operator: str, version: str):
from packaging import version as pversion
torch_version = pversion.parse(torch.__version__)
target_version = pversion.parse(version)
# print(f"torch_version: {torch_version}, target: torch{operator}{target_version}")
if operator == ">":
return torch_version > target_version
elif operator == ">=":
return torch_version >= target_version
elif operator == "==":
return torch_version == target_version
elif operator == "<=":
return torch_version <= target_version
elif operator == "<":
return torch_version < target_version
return False
|