installs an older version of Triton if CUDA < 11.4

setup.py

@@ -12,7 +12,18 @@ def read_version(fname="whisper/version.py"):
 
 requirements = []
 if sys.platform.startswith("linux"):
- requirements.append("triton>=2.0.0.dev20221202")
+ triton_requirement = "triton>=2.0.0.dev20221202"
+ try:
+ import re
+ import subprocess
+ version_line = subprocess.check_output(["nvcc", "--version"]).strip().split(b"\n")[-1]
+ major, minor = re.findall(rb"cuda_([\d]+)\.([\d]+)", version_line)[0]
+ if (int(major), int(minor)) < (11, 4):
+ # the last version supporting CUDA < 11.4
+ triton_requirement = "triton==2.0.0.dev20221011"
+ except (IndexError, OSError, subprocess.SubprocessError):
+ pass
+ requirements.append(triton_requirement)
 
 setup(
  name="openai-whisper",