add package name for cuda 12

microsoft · Nov 20, 2023 · 87d9500 · 87d9500
1 parent a24010c
commit 87d9500
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/setup.py b/setup.py
@@ -536,6 +536,9 @@ def finalize_options(self):
                 # Training ROCM package for ADO feeds is called onnxruntime-training-rocm
                 package_name = "onnxruntime-training-rocm"
 
+            if cuda_version and cuda_version.split(".")[0] == "12":
+                package_name = "onnxruntime-training-cuda12"
+
 if package_name == "onnxruntime-tvm":
     packages += ["onnxruntime.providers.tvm"]