diff --git a/docker/Dockerfile b/docker/Dockerfile
index aa9f3a4e4cf..1901d4c27a1 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -26,6 +26,15 @@ WORKDIR /sgl-workspace
 ARG CUDA_VERSION
 RUN python3 -m pip install --upgrade pip setuptools wheel html5lib six \
     && git clone --depth=1 https://github.com/sgl-project/sglang.git \
+    && if [ "$CUDA_VERSION" = "12.1.1" ]; then \
+         python3 -m pip install torch --index-url https://download.pytorch.org/whl/cu121; \
+       elif [ "$CUDA_VERSION" = "12.4.1" ]; then \
+         python3 -m pip install torch --index-url https://download.pytorch.org/whl/cu124; \
+       elif [ "$CUDA_VERSION" = "11.8.0" ]; then \
+         python3 -m pip install torch --index-url https://download.pytorch.org/whl/cu118; \
+       else \
+         echo "Unsupported CUDA version: $CUDA_VERSION" && exit 1; \
+       fi \
     && cd sglang \
     && if [ "$BUILD_TYPE" = "srt" ]; then \
          if [ "$CUDA_VERSION" = "12.1.1" ]; then \