diff --git a/docker/Dockerfile b/docker/Dockerfile
index 69d1c974..377b5164 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -63,6 +63,7 @@ WORKDIR /app/ai-toolkit
 
 # Install Python dependencies
 RUN pip install --no-cache-dir -r requirements.txt && \
+    pip install --pre --no-cache-dir torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu128 --force && \
     pip install setuptools==69.5.1 --no-cache-dir
 
 # Build UI