github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[setup] fixed inconsistent version meta (#2578)

pull/2585/head
Frank Lee 2023-02-06 13:48:20 +08:00 committed by GitHub
parent 8af5a0799b
commit fba08743a8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 16 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
setup.py
View File

@ -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