From d798ff739d9aedf616621b6b1153455cc1b10239 Mon Sep 17 00:00:00 2001 From: FrankLeeeee Date: Mon, 6 Feb 2023 09:36:49 +0800 Subject: [PATCH] [setup] fixed inconsistent version meta --- setup.py | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/setup.py b/setup.py index b9cd9e5e4714..7cfbbe9b19a4 100644 --- a/setup.py +++ b/setup.py @@ -117,14 +117,26 @@ def get_version(): with open(version_txt_path) as f: 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}' # write version into version.py with open(version_py_path, 'w') as f: f.write(f"__version__ = '{version}'\n") + if build_cuda_ext: + torch_version = '.'.join(torch.__version__.split('.')[:2]) + cuda_version = '.'.join(get_cuda_bare_metal_version(CUDA_HOME)[1:]) + else: + torch_version = None + cuda_version = None + + if torch_version: + f.write(f'torch = "{torch_version}"\n') + else: + f.write('torch = None\n') + + if cuda_version: + f.write(f'cuda = "{cuda_version}"\n') + else: + f.write('cuda = None\n') return version