Propagate ROCM version to onnxruntime wheel package (microsoft#8247) (microsoft#8250)

Co-authored-by: Thiago Crepaldi <[email protected]>

Author: harshithapv

setup.py

@@ -424,7 +424,7 @@ def get_torch_version():
     install_requires = f.read().splitlines()
 
 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
 
@@ -446,11 +446,10 @@ def save_build_and_package_info(package_name, version_number, cuda_version):
                         "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(