saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-05-14 20:48:00 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Propagate ROCM version to onnxruntime wheel package (#8247)

Browse source
This commit is contained in:
Thiago Crepaldi 2021-06-30 13:52:22 -07:00 committed by GitHub
parent 665ecdf9ce
commit 97f1eea2ea
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
setup.py
View file

@ -428,7 +428,7 @@ with open(requirements_path) as f:
if enable_training:
def save_build_and_package_info(package_name, version_number, cuda_version):
def save_build_and_package_info(package_name, version_number, cuda_version, rocm_version):
sys.path.append(path.join(path.dirname(__file__), 'onnxruntime', 'python'))
from onnxruntime_collect_build_info import find_cudart_versions
@ -450,11 +450,10 @@ if enable_training:
"did not find any cudart library"
if not cudart_versions or len(cudart_versions) == 0
else "found multiple cudart libraries")
else:
# TODO: rocm
pass
elif rocm_version:
f.write("rocm_version = '{}'\n".format(rocm_version))
save_build_and_package_info(package_name, version_number, cuda_version)
save_build_and_package_info(package_name, version_number, cuda_version, rocm_version)
# Setup
setup(