From daf59ff72ef992312b59348939053de711cea38b Mon Sep 17 00:00:00 2001
From: ver217 <lhx0217@gmail.com>
Date: Wed, 27 Apr 2022 15:26:12 +0800
Subject: [PATCH] [setup] add local version label (#890)

---
 setup.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index a6de9cc49..d995ef4c4 100644
--- 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: