diff --git a/official-templates/better-base-rocm/Dockerfile b/official-templates/better-base-rocm/Dockerfile
index c170d66..6ec5d38 100644
--- a/official-templates/better-base-rocm/Dockerfile
+++ b/official-templates/better-base-rocm/Dockerfile
@@ -29,10 +29,13 @@ 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 curl -fsSL https://code-server.dev/install.sh | sh
+RUN for i in 1 2 3; do \
     code-server --install-extension ms-python.python && \
-    code-server --install-extension /tmp/ms-toolsai.jupyter-2024.7.0.vsix && \
     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
 RUN rm /tmp/ms-toolsai.jupyter-2024.7.0.vsix
 
 # Pre-install Jupyter kernel and OhMyRunPod