fix: install torch

blib-la · Oct 6, 2024 · 97b2c2e · 97b2c2e
1 parent 894eaca
commit 97b2c2e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/Dockerfile b/Dockerfile
@@ -5,7 +5,8 @@ FROM runpod/base:0.6.1-cuda12.2.0
 RUN apt-get update && apt-get install -y ffmpeg aria2 git unzip
 
 # Install Python dependencies
-RUN pip install --no-cache-dir opencv-python imageio imageio-ffmpeg ffmpeg-python av runpod \
+RUN pip install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121 \
+    opencv-python imageio imageio-ffmpeg ffmpeg-python av runpod \
     xformers==0.0.25 torchsde==0.2.6 einops==0.8.0 diffusers==0.28.0 transformers==4.41.2 accelerate==0.30.1
 
 # Clone the ComfyUI repository