fix #2528 (#2541)

sgl-project · Dec 21, 2024 · 4e1e3cf · 4e1e3cf
1 parent 8f4d04e
commit 4e1e3cf
Showing 1 changed file with 9 additions and 0 deletions.
diff --git 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 \