diff --git a/maxtext_gpu_dependencies.Dockerfile b/maxtext_gpu_dependencies.Dockerfile
index ef718a44..b09cf520 100644
--- a/maxtext_gpu_dependencies.Dockerfile
+++ b/maxtext_gpu_dependencies.Dockerfile
@@ -38,12 +38,13 @@ RUN mkdir -p /deps
 # Set the working directory in the container
 WORKDIR /deps
 
-# Copy all files from local workspace into docker container
-COPY . .
-RUN ls .
+# Copy setup files and dependency files separately for better caching
+COPY setup.sh ./
+COPY constraints_gpu.txt requirements.txt requirements_with_jax_stable_stack.txt ./
 
+# Install dependencies - these steps are cached unless the copied files change
 RUN echo "Running command: bash setup.sh MODE=$ENV_MODE JAX_VERSION=$ENV_JAX_VERSION DEVICE=${ENV_DEVICE}"
 RUN --mount=type=cache,target=/root/.cache/pip bash setup.sh MODE=${ENV_MODE} JAX_VERSION=${ENV_JAX_VERSION} DEVICE=${ENV_DEVICE}
 
-
-WORKDIR /deps
+# Now copy the remaining code (source files that may change frequently)
+COPY . .