diff --git a/official-templates/better-base-rocm/Dockerfile b/official-templates/better-base-rocm/Dockerfile
index 6ec5d38..242fcdf 100644
--- a/official-templates/better-base-rocm/Dockerfile
+++ b/official-templates/better-base-rocm/Dockerfile
@@ -29,13 +29,11 @@ RUN pip install --upgrade --no-cache-dir pip
 COPY ms-toolsai.jupyter-2024.7.0.vsix /tmp/ms-toolsai.jupyter-2024.7.0.vsix
 
 # Install code-server and extensions
-RUN curl -fsSL https://code-server.dev/install.sh | sh
-RUN for i in 1 2 3; do \
+RUN curl -fsSL https://code-server.dev/install.sh | sh && \
     code-server --install-extension ms-python.python && \
-    code-server --install-extension ms-toolsai.vscode-jupyter-powertoys
-      break || sleep 10; \
-    done
-RUN code-server --install-extension /tmp/ms-toolsai.jupyter-2024.7.0.vsix
+    code-server --install-extension ms-toolsai.vscode-jupyter-powertoys && \
+    code-server --install-extension /tmp/ms-toolsai.jupyter-2024.7.0.vsix
+    
 RUN rm /tmp/ms-toolsai.jupyter-2024.7.0.vsix
 
 # Pre-install Jupyter kernel and OhMyRunPod