diff --git a/tools/ci_build/github/azure-pipelines/linux-rocm-ci-pipeline.yml b/tools/ci_build/github/azure-pipelines/linux-rocm-ci-pipeline.yml
index 60cd9cf1d8..7a27b3a340 100644
--- a/tools/ci_build/github/azure-pipelines/linux-rocm-ci-pipeline.yml
+++ b/tools/ci_build/github/azure-pipelines/linux-rocm-ci-pipeline.yml
@@ -94,7 +94,6 @@ jobs:
           /bin/bash -c "
             set -ex; \
             ccache -s; \
-            source /ort/env/bin/activate; \
             env; \
             python tools/ci_build/build.py \
               --config Release \
@@ -184,7 +183,6 @@ jobs:
           /bin/bash -c "
             set -ex; \
             xargs -a /build/Release/perms.txt chmod a+x; \
-            source /ort/env/bin/activate; \
             python /onnxruntime_src/tools/ci_build/build.py \
               --config Release \
               --cmake_extra_defines \
@@ -202,7 +200,7 @@ jobs:
               --parallel \
               --build_wheel \
               --skip_submodule_sync \
-              --test --enable_onnx_tests --enable_transformers_tool_test \
+              --test --enable_onnx_tests \
               --cmake_path /usr/bin/cmake --ctest_path /usr/bin/ctest"
       workingDirectory: $(Build.SourcesDirectory)
     displayName: 'Run onnxruntime unit tests'
@@ -220,6 +218,7 @@ jobs:
           --user $UID:$(id -g $USER) \
           --volume $(Build.SourcesDirectory):/onnxruntime_src \
           --volume $(Build.BinariesDirectory):/build \
+          --workdir /build/Release \
           -e OPENBLAS_NUM_THREADS=1 \
           -e OPENMP_NUM_THREADS=1 \
           -e MKL_NUM_THREADS=1 \
@@ -227,11 +226,9 @@ jobs:
           -e KERNEL_EXPLORER_BATCHED_GEMM_MAX_BATCH_SIZE=8 \
           -e KERNEL_EXPLORER_TEST_USE_CUPY=0 \
           -e CUPY_CACHE_DIR=/build/Release \
-          --workdir /ort \
           onnxruntime-rocm-cibuild-rocm$(RocmVersion) \
           /bin/bash -c "
             set -ex; \
-            source /ort/env/bin/activate; \
             python --version; \
             pytest /onnxruntime_src/onnxruntime/python/tools/kernel_explorer/ -n 4 --reruns 1 --durations=100"
       workingDirectory: $(Build.SourcesDirectory)
diff --git a/tools/ci_build/github/linux/docker/Dockerfile.ubuntu_rocm b/tools/ci_build/github/linux/docker/Dockerfile.ubuntu_rocm
index b1941ceceb..ca0b915f20 100644
--- a/tools/ci_build/github/linux/docker/Dockerfile.ubuntu_rocm
+++ b/tools/ci_build/github/linux/docker/Dockerfile.ubuntu_rocm
@@ -15,8 +15,9 @@ LABEL PYTHON_VERSION="${PYTHON_VERSION}"
 ENV DEBIAN_FRONTEND=noninteractive \
     LC_ALL=C.UTF-8 \
     LANG=C.UTF-8 \
-    LD_LIBRARY_PATH=/opt/rocm/lib:/usr/lib/x86_64-linux-gnu \
-    PATH=/opt/rocm/bin:/usr/bin:/bin:/usr/sbin:/usr/local/bin
+    VIRTUAL_ENV=/opt/venv \
+    LD_LIBRARY_PATH=/opt/rocm/lib:/usr/lib/x86_64-linux-gnu::/opt/amdgpu/lib/x86_64-linux-gnu \
+    PATH=/opt/venv/bin:/opt/rocm/bin:/usr/bin:/bin:/usr/sbin:/usr/local/bin
 
 # Set default shell
 SHELL ["/bin/bash", "-c"]
@@ -80,13 +81,13 @@ RUN mkdir -p /tmp/ccache && \
     rm -rf /tmp/ccache
 
 # Set up virtual environment for Python and install dependencies
-WORKDIR /ort
 COPY scripts/requirements.txt /ort/
-RUN python3 -m venv /ort/env && \
-    source /ort/env/bin/activate && \
+RUN python3 -m venv ${VIRTUAL_ENV} && \
+    source ${VIRTUAL_ENV}/bin/activate && \
     pip install --upgrade pip setuptools wheel && \
     pip install -r /ort/requirements.txt && \
     pip install ml_dtypes pytest-xdist pytest-rerunfailures scipy
 
+
 # Default command
 CMD ["/bin/bash"]