Provide pre-built wheels with CUDA support. (#1143)

This commit is contained in:
Fangjun Kuang
2024-07-17 22:59:13 +08:00
committed by GitHub
parent 803c02db0a
commit 4198d9a166
4 changed files with 233 additions and 2 deletions

View File

@@ -1,5 +1,6 @@
#!/usr/bin/env python3
import os
import re
from pathlib import Path
@@ -26,6 +27,14 @@ def get_package_version():
match = re.search(r"set\(SHERPA_ONNX_VERSION (.*)\)", content)
latest_version = match.group(1).strip('"')
cmake_args = os.environ.get("SHERPA_ONNX_CMAKE_ARGS", "")
extra_version = ""
if "-DSHERPA_ONNX_ENABLE_GPU=ON" in cmake_args:
extra_version = "+cuda"
latest_version += extra_version
return latest_version