diff --git a/docker/base-dev-x.dockerfile b/docker/base-dev-x.dockerfile
index 5c0a3192e..5e1d5963a 100644
--- a/docker/base-dev-x.dockerfile
+++ b/docker/base-dev-x.dockerfile
@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=ghcr.io/microsoft/ark/ark:base-cuda12.1
+ARG BASE_IMAGE
 FROM ${BASE_IMAGE}
 
 LABEL maintainer="ARK"
@@ -10,10 +10,11 @@ ENV ARK_SRC_DIR="/tmp/ark" \
 ADD . ${ARK_SRC_DIR}
 WORKDIR ${ARK_SRC_DIR}
 
-# Install Lcov
+# Install Lcov and clang-format
 RUN apt-get update && \
     apt-get install -y --no-install-recommends \
         lcov \
+        clang-format \
         && \
     apt-get autoremove && \
     apt-get clean && \
@@ -27,6 +28,14 @@ RUN curl -L ${CMAKE_URL} -o ${CMAKE_HOME}.tar.gz && \
     rm -rf ${CMAKE_HOME}.tar.gz
 ENV PATH="/usr/local/cmake-${CMAKE_VERSION}-linux-x86_64/bin:${PATH}"
 
+# Install Python dependencies
+ADD . /tmp/ark
+WORKDIR /tmp/ark
+RUN python3 -m pip install --no-cache-dir -r requirements.txt
+
+# Install PyTorch and black
+RUN python3 -m pip install --no-cache-dir torch black
+
 # Set PATH
 RUN echo PATH="${PATH}" > /etc/environment
 
diff --git a/docker/base-rocm5.6.dockerfile b/docker/base-rocm5.6.dockerfile
index 9ba35e4a0..f92abc4ac 100644
--- a/docker/base-rocm5.6.dockerfile
+++ b/docker/base-rocm5.6.dockerfile
@@ -1,6 +1,6 @@
 # Temporal Dockerfile for building ARK base image for ROCm 5.6
 
-ARG BASE_IMAGE=rocm/dev-ubuntu-20.04:5.6.1-complete
+ARG BASE_IMAGE
 FROM ${BASE_IMAGE}
 
 LABEL maintainer="ARK"
diff --git a/docker/base-x.dockerfile b/docker/base-x.dockerfile
index c5ff03afe..c63ccc6b1 100644
--- a/docker/base-x.dockerfile
+++ b/docker/base-x.dockerfile
@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=nvidia/cuda:12.1.1-devel-ubuntu20.04
+ARG BASE_IMAGE
 FROM ${BASE_IMAGE}
 
 LABEL maintainer="ARK"