[setup] add local version label (#890)

daf59ff7 · ver217 · GitHub · 1d0aba41 · daf59ff7
Unverified Commit daf59ff7 authored Apr 27, 2022 by ver217 Committed by GitHub Apr 27, 2022
Show whitespace changes
Inline Side-by-side

Showing with 6 additions and 1 deletion

setup.py setup.py +6 -1

No files found.
--- a/setup.py
+++ b/setup.py
@@ -94,7 +94,12 @@ def fetch_readme():

 def get_version():
    with open('version.txt') as f:
-        return f.read().strip()
+        version = f.read().strip()
+        if build_cuda_ext:
+            torch_version = ''.join(torch.__version__.split('.')[:2])
+            cuda_version = ''.join(get_cuda_bare_metal_version(CUDA_HOME)[1:])
+            version += f'+torch{torch_version}cu{cuda_version}'
+        return version


 if build_cuda_ext: