diff --git a/docker/Dockerfile b/docker/Dockerfile
index e07b781e..76835bfc 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -37,6 +37,13 @@ RUN curl -sL https://deb.nodesource.com/setup_23.x -o nodesource_setup.sh && \
 
 WORKDIR /app
 
+# Set aliases for python and pip
+RUN echo 'alias python=python3' >> /etc/bash.bashrc && \
+    echo 'alias pip=pip3' >> /etc/bash.bashrc
+
+# install pytorch before cache bust to avoid redownloading pytorch
+RUN pip install --no-cache-dir torch==2.6.0 torchvision==0.21.0 --index-url https://download.pytorch.org/whl/cu126
+
 # Fix cache busting by moving CACHEBUST to right before git clone
 ARG CACHEBUST=1234
 RUN echo "Cache bust: ${CACHEBUST}" && \
@@ -46,13 +53,8 @@ RUN echo "Cache bust: ${CACHEBUST}" && \
 
 WORKDIR /app/ai-toolkit
 
-# Set aliases for python and pip
-RUN echo 'alias python=python3' >> /etc/bash.bashrc && \
-    echo 'alias pip=pip3' >> /etc/bash.bashrc
-
 # Install Python dependencies
-RUN pip install --no-cache-dir torch==2.6.0 torchvision==0.21.0 --index-url https://download.pytorch.org/whl/cu126 && \
-    pip install --no-cache-dir -r requirements.txt
+RUN pip install --no-cache-dir -r requirements.txt
 
 # Build UI
 WORKDIR /app/ai-toolkit/ui