diff --git a/engine/Dockerfile.nvidia b/engine/Dockerfile.nvidia
index 1064f07..e9f602c 100644
--- a/engine/Dockerfile.nvidia
+++ b/engine/Dockerfile.nvidia
@@ -20,8 +20,8 @@ COPY VERSION ./
 
 RUN uv venv .venv && \
     . .venv/bin/activate && \
-    uv pip install -e . && \
-    uv pip install --no-deps onnxruntime-gpu
+    UV_HTTP_TIMEOUT=600 uv pip install torch torchvision --index-url https://download.pytorch.org/whl/cu130 && \
+    uv pip install -e .
 
 ENV PATH="/app/.venv/bin:$PATH"
 ENV VIRTUAL_ENV="/app/.venv"
diff --git a/engine/Dockerfile.rocm b/engine/Dockerfile.rocm
index adc4008..4a45274 100644
--- a/engine/Dockerfile.rocm
+++ b/engine/Dockerfile.rocm
@@ -18,10 +18,15 @@ COPY kb/ kb/
 COPY main.py ./
 COPY VERSION ./
 
+COPY torch-*.whl /tmp/
+
 RUN uv venv .venv && \
     . .venv/bin/activate && \
+    uv pip install /tmp/torch-*.whl --extra-index-url https://download.pytorch.org/whl/rocm6.4 && \
+    uv pip install torchvision --index-url https://download.pytorch.org/whl/rocm6.4 && \
     uv pip install -e . && \
-    uv pip install --no-deps onnxruntime-rocm
+    rm /tmp/torch-*.whl && \
+    rm -rf .venv/lib/python3.12/site-packages/nvidia
 
 # Stage 2: Runtime — minimal ROCm runtime libs only
 FROM ubuntu:24.04