diff --git a/tools/ci_build/github/azure-pipelines/orttraining-py-packaging-pipeline-cuda.yml b/tools/ci_build/github/azure-pipelines/orttraining-py-packaging-pipeline-cuda.yml
index 539a61c021..2c65432471 100644
--- a/tools/ci_build/github/azure-pipelines/orttraining-py-packaging-pipeline-cuda.yml
+++ b/tools/ci_build/github/azure-pipelines/orttraining-py-packaging-pipeline-cuda.yml
@@ -33,3 +33,4 @@ stages:
     debug_build: false
     SpecificArtifact: ${{ parameters.SpecificArtifact }}
     BuildId: ${{ parameters.BuildId }}
+    build_pool_name: 'onnxruntime-Ubuntu2204-AMD-CPU'
diff --git a/tools/ci_build/github/azure-pipelines/orttraining-py-packaging-pipeline-cuda12.yml b/tools/ci_build/github/azure-pipelines/orttraining-py-packaging-pipeline-cuda12.yml
index 86dce7ae46..6e02714f28 100644
--- a/tools/ci_build/github/azure-pipelines/orttraining-py-packaging-pipeline-cuda12.yml
+++ b/tools/ci_build/github/azure-pipelines/orttraining-py-packaging-pipeline-cuda12.yml
@@ -20,3 +20,4 @@ stages:
     agent_pool: Onnxruntime-Linux-GPU
     upload_wheel: 'yes'
     debug_build: false
+    build_pool_name: 'onnxruntime-Linux-GPU'
diff --git a/tools/ci_build/github/azure-pipelines/templates/py-packaging-training-cuda-stage-steps.yml b/tools/ci_build/github/azure-pipelines/templates/py-packaging-training-cuda-stage-steps.yml
index 08f001efb6..f6b36733eb 100644
--- a/tools/ci_build/github/azure-pipelines/templates/py-packaging-training-cuda-stage-steps.yml
+++ b/tools/ci_build/github/azure-pipelines/templates/py-packaging-training-cuda-stage-steps.yml
@@ -11,6 +11,7 @@ parameters:
   stage_name: ''
   SpecificArtifact: false
   BuildId: '0'
+  build_pool_name: ''
 
 stages:
   - stage: Build_${{ parameters.stage_name }}
@@ -30,12 +31,12 @@ stages:
       - name: PythonVersion
         value: ${{ parameters.python_version }}
       - name: Repository
-        value: onnxruntimetraininggpubuild_${{ parameters.python_version }}
+        value: onnxruntimetraininggpubuild_cu${{ replace(parameters.cuda_version, '.', '') }}_py${{ replace(parameters.python_version, '.', '') }}
     dependsOn: []
 
     jobs:
     - job: Build
-      pool: onnxruntime-Ubuntu2204-AMD-CPU
+      pool: ${{ parameters.build_pool_name }}
       timeoutInMinutes: 180
       steps:
         - task: mspremier.PostBuildCleanup.PostBuildCleanup-task.PostBuildCleanup@3
@@ -84,7 +85,6 @@ stages:
                 --volume $(Build.BinariesDirectory):/build \
                 --volume /data/models:/build/models:ro \
                 --volume $HOME/.onnx:/home/onnxruntimedev/.onnx \
-                -e NVIDIA_VISIBLE_DEVICES=all \
                 -e NIGHTLY_BUILD \
                 -e DEFAULT_TRAINING_PACKAGE_DEVICE \
                 -e BUILD_BUILDNUMBER \
@@ -137,7 +137,7 @@ stages:
       - name: PythonVersion
         value: ${{ parameters.python_version }}
       - name: Repository
-        value: onnxruntimetraininggpubuild_${{ parameters.python_version }}
+        value: onnxruntimetraininggpubuild_cu${{ replace(parameters.cuda_version, '.', '') }}_py${{ replace(parameters.python_version, '.', '') }}
       - name: UploadWheel
         value: ${{ parameters.upload_wheel }}
     dependsOn: Build_${{ parameters.stage_name }}
diff --git a/tools/ci_build/github/azure-pipelines/templates/py-packaging-training-cuda-stage.yml b/tools/ci_build/github/azure-pipelines/templates/py-packaging-training-cuda-stage.yml
index f7ecc3cf84..a1f326ebaa 100644
--- a/tools/ci_build/github/azure-pipelines/templates/py-packaging-training-cuda-stage.yml
+++ b/tools/ci_build/github/azure-pipelines/templates/py-packaging-training-cuda-stage.yml
@@ -57,6 +57,11 @@ parameters:
   type: string
   default: '0'
 
+- name: build_pool_name
+  displayName: >
+    build_pool_name.
+  type: string
+
 - name: PythonVersionList
   displayName: Python Version List
   type: object
@@ -86,3 +91,4 @@ stages:
       python_version: ${{ python_version.version }}
       SpecificArtifact: ${{ parameters.SpecificArtifact }}
       BuildId: ${{ parameters.BuildId }}
+      build_pool_name: ${{ parameters.build_pool_name }}