From 97f1eea2ea8e068dd29bd3fefa6043bc067f7eb4 Mon Sep 17 00:00:00 2001 From: Thiago Crepaldi Date: Wed, 30 Jun 2021 13:52:22 -0700 Subject: [PATCH] Propagate ROCM version to onnxruntime wheel package (#8247) --- setup.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/setup.py b/setup.py index a3bf000d85..18f5fac3e0 100644 --- a/setup.py +++ b/setup.py @@ -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(