diff --git a/tools/ci_build/github/azure-pipelines/orttraining-pai-ci-pipeline.yml b/tools/ci_build/github/azure-pipelines/orttraining-pai-ci-pipeline.yml
index f4d10699c0..2dd9ab6e0d 100644
--- a/tools/ci_build/github/azure-pipelines/orttraining-pai-ci-pipeline.yml
+++ b/tools/ci_build/github/azure-pipelines/orttraining-pai-ci-pipeline.yml
@@ -203,18 +203,6 @@ jobs:
     condition: and(succeededOrFailed(), eq(variables.onnxruntimeBuildSucceeded, 'true'))
 
 
-  - bash: tools/ci_build/github/linux/docker/scripts/training/azure_scale_set_vm_mount_test_data.sh -p $(orttrainingtestdatascus-storage-key) -s "//orttrainingtestdatascus.file.core.windows.net/mnist" -d "$(Build.SourcesDirectory)/mnist"
-    displayName: 'Mount MNIST'
-    condition: and(succeededOrFailed(), eq(variables.onnxruntimeBuildSucceeded, 'true'))
-
-  - bash: tools/ci_build/github/linux/docker/scripts/training/azure_scale_set_vm_mount_test_data.sh -p $(orttrainingtestdatascus-storage-key) -s "//orttrainingtestdatascus.file.core.windows.net/bert-data" -d "$(Build.SourcesDirectory)/bert_data"
-    displayName: 'Mount bert-data'
-    condition: and(succeededOrFailed(), eq(variables.onnxruntimeBuildSucceeded, 'true'))
-
-  - bash: tools/ci_build/github/linux/docker/scripts/training/azure_scale_set_vm_mount_test_data.sh -p $(orttrainingtestdatascus-storage-key) -s "//orttrainingtestdatascus.file.core.windows.net/hf-models-cache" -d "$(Build.SourcesDirectory)/hf_models_cache"
-    displayName: 'Mount hf-models-cache'
-    condition: and(succeededOrFailed(), eq(variables.onnxruntimeBuildSucceeded, 'true'))
-
   # Entry point for all ORTModule tests
   # The onnxruntime folder is deleted in the build directory
   # to enforce use of the onnxruntime wheel
@@ -227,22 +215,24 @@ jobs:
         whlfilename=$(basename ${files[0]})
         echo $whlfilename
         docker run --rm \
+          -e HIP_VISIBLE_DEVICES \
           --security-opt seccomp=unconfined \
           --shm-size=1024m \
           --device=/dev/kfd \
-          --device=/dev/dri/renderD$DRIVER_RENDER \
+          --device=/dev/dri \
+          --privileged \
           --group-add $(video) \
           --group-add $(render) \
           --user onnxruntimedev \
           --volume $(Build.BinariesDirectory):/build \
-          --volume $(Build.SourcesDirectory)/mnist:/mnist \
-          --volume $(Build.SourcesDirectory)/bert_data:/bert_data \
-          --volume $(Build.SourcesDirectory)/hf_models_cache:/hf_models_cache \
+          --volume $(Build.SourcesDirectory):/onnxruntime_src \
           --workdir /build/$(BuildConfig) \
           onnxruntimetrainingrocm-cibuild-rocm$(RocmVersion) \
             /bin/bash -c "
               set -ex; \
               unset PYTHONPATH; \
+              /onnxruntime_src/tools/ci_build/github/linux/docker/scripts/training/azure_scale_set_vm_mount_test_data.sh -p $(orttrainingtestdatascus-storage-key) -s "//orttrainingtestdatascus.file.core.windows.net/mnist" -d "/mnist"; \
+              /onnxruntime_src/tools/ci_build/github/linux/docker/scripts/training/azure_scale_set_vm_mount_test_data.sh -p $(orttrainingtestdatascus-storage-key) -s "//orttrainingtestdatascus.file.core.windows.net/bert-data" -d "/bert_data"; \
               pip install /build/$(BuildConfig)/dist/$whlfilename; \
               python -m onnxruntime.training.ortmodule.torch_cpp_extensions.install; \
               python orttraining_ortmodule_tests.py \
@@ -252,22 +242,6 @@ jobs:
     displayName: 'Run orttraining_ortmodule_tests.py'
     condition: and(succeededOrFailed(), eq(variables.onnxruntimeBuildSucceeded, 'true'))
 
-  - task: CmdLine@2
-    inputs:
-      script: |-
-        if [ -d $(Build.SourcesDirectory)/mnist ]; then
-          sudo umount $(Build.SourcesDirectory)/mnist
-        fi
-        if [ -d $(Build.SourcesDirectory)/bert_data ]; then
-          sudo umount $(Build.SourcesDirectory)/bert_data
-        fi
-        if [ -d $(Build.SourcesDirectory)/hf_models_cache ]; then
-          sudo umount $(Build.SourcesDirectory)/hf_models_cache
-        fi
-      workingDirectory: $(Build.SourcesDirectory)
-    displayName: 'Umount dataset'
-    condition: always()
-
   - template: templates/component-governance-component-detection-steps.yml
     parameters :
       condition : 'succeeded'
diff --git a/tools/ci_build/github/pai/rocm-ci-pipeline-env.Dockerfile b/tools/ci_build/github/pai/rocm-ci-pipeline-env.Dockerfile
index 42617d6b4c..23cd963056 100644
--- a/tools/ci_build/github/pai/rocm-ci-pipeline-env.Dockerfile
+++ b/tools/ci_build/github/pai/rocm-ci-pipeline-env.Dockerfile
@@ -3,6 +3,8 @@ FROM rocm/pytorch:rocm5.5_ubuntu20.04_py3.8_pytorch_1.13.1
 ARG BUILD_UID=1001
 ARG BUILD_USER=onnxruntimedev
 RUN adduser --uid $BUILD_UID $BUILD_USER
+RUN echo "$BUILD_USER ALL=(ALL) NOPASSWD:ALL" >> /etc/sudoers.d/$BUILD_USER
+
 WORKDIR /home/$BUILD_USER
 
 # from rocm/pytorch's image, work around ucx's dlopen replacement conflicting with shared provider
@@ -25,6 +27,8 @@ RUN mkdir -p /tmp/ccache && \
     cp /tmp/ccache/ccache /usr/bin && \
     rm -rf /tmp/ccache
 
+RUN apt-get update && apt-get install  -y cifs-utils
+
 USER $BUILD_USER
 
 # rocm-ci branch contains instrumentation needed for loss curves and perf