From 87d95003163ab91e1a42004d4fa2080ea46b3c1a Mon Sep 17 00:00:00 2001 From: Abhishek Jindal Date: Mon, 20 Nov 2023 13:47:59 -0800 Subject: [PATCH] add package name for cuda 12 --- setup.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/setup.py b/setup.py index 1c04433c9a7ca..4df48239c8cbd 100644 --- 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"]