diff --git a/docker/Dockerfile b/docker/Dockerfile
index 08957cfe96..bc738ad90d 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -94,10 +94,10 @@ RUN \
     elif echo "${BASE_IMAGE}" | grep -q "rocm/"; then \
         # Install ROCm version specific binary when ROCm version is specified as a build arg
         if [ "$USE_ROCM_VERSION" ]; then \
-            python ./ts_scripts/install_dependencies.py --rocm $USE_ROCM_VERSION\
+            python ./ts_scripts/install_dependencies.py --rocm $USE_ROCM_VERSION;\
         # Install the binary with the latest CPU image on a ROCm base image
         else \
-            python ./ts_scripts/install_dependencies.py;\
+            python ./ts_scripts/install_dependencies.py; \
         fi; \
     # Install the CPU binary
     else \
@@ -192,7 +192,6 @@ RUN --mount=type=cache,target=/var/cache/apt \
     && rm -rf /var/lib/apt/lists/* \
     && cd /tmp
 
-
 COPY --from=compile-image /home/venv /home/venv
 
 ENV PATH="/home/venv/bin:$PATH"
@@ -211,6 +210,7 @@ FROM ${BASE_IMAGE} as dev-image
 ARG PYTHON_VERSION
 ARG BRANCH_NAME
 ARG BUILD_FROM_SRC
+ARG USE_ROCM_VERSION
 ARG LOCAL_CHANGES
 ARG BUILD_WITH_IPEX
 ARG IPEX_VERSION=1.11.0
diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm
index f9ff5dce79..ccefd7b944 100644
--- a/docker/Dockerfile.rocm
+++ b/docker/Dockerfile.rocm
@@ -277,16 +277,12 @@ WORKDIR "/serve"
 
 RUN \
     if [ "$USE_ROCM_VERSION" ]; then \
-        python ts_scripts/install_dependencies.py --environment=dev --rocm $USE_ROCM_VERSION \
-        && cd /opt/rocm/share/amd_smi \
-        && python -m pip install . ; \
-    # Install the binary with the latest CPU image on a ROCm base image
-    else \
-        python ts_scripts/install_dependencies.py --environment=dev;\
+        python -m pip install /opt/rocm/share/amd_smi ; \
     fi;
 
 RUN python -m pip install -U pip setuptools \
     && python -m pip install --no-cache-dir -r requirements/developer.txt \
+    && python ts_scripts/install_dependencies.py --environment dev \
     && python ts_scripts/install_from_src.py --environment=dev \
     && useradd -m model-server \
     && mkdir -p /home/model-server/tmp \
diff --git a/docker/build_image.sh b/docker/build_image.sh
index f998150a66..a6ddac2fee 100755
--- a/docker/build_image.sh
+++ b/docker/build_image.sh
@@ -269,7 +269,7 @@ if [ "${BUILD_TYPE}" == "production" ]; then
   fi
 elif [ "${BUILD_TYPE}" == "ci" ];
 then
-  DOCKER_BUILDKIT=1 docker build --file Dockerfile --build-arg BASE_IMAGE="${BASE_IMAGE}" --build-arg USE_CUDA_VERSION="${CUDA_VERSION}"  --build-arg PYTHON_VERSION="${PYTHON_VERSION}"\
+  DOCKER_BUILDKIT=1 docker build --file Dockerfile --build-arg BASE_IMAGE="${BASE_IMAGE}" --build-arg USE_CUDA_VERSION="${CUDA_VERSION}" --build-arg USE_ROCM_VERSION="${ROCM_VERSION}" --build-arg PYTHON_VERSION="${PYTHON_VERSION}"\
   --build-arg BUILD_NIGHTLY="${BUILD_NIGHTLY}" --build-arg BRANCH_NAME="${BRANCH_NAME}" --build-arg REPO_URL="${REPO_URL}" --build-arg BUILD_FROM_SRC="${BUILD_FROM_SRC}"\
   --build-arg LOCAL_CHANGES="${LOCAL_CHANGES}" -t "${DOCKER_TAG}" --target ci-image  ../
 else