Install correct version of Torch

ashleykleynhans · Oct 20, 2023 · 1d71c18 · 1d71c18
1 parent bfc8345
commit 1d71c18
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/Dockerfile b/Dockerfile
@@ -63,6 +63,12 @@ FROM base as setup
 # Create and use the Python venv
 RUN python3 -m venv /venv
 
+# Install Torch
+RUN source /venv/bin/activate && \
+    pip3 install --no-cache-dir torch==2.0.1 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 && \
+    pip3 install --no-cache-dir xformers==0.0.22 && \
+    deactivate
+
 # Clone the git repo of LLaVA and set version
 WORKDIR /
 RUN git clone https://github.com/ashleykleynhans/LLaVA.git

diff --git a/README.md b/README.md
@@ -7,7 +7,7 @@
 * Python 3.10.12
 * [LLaVA](
   https://github.com/haotian-liu/llava) v1.1.1
-* Torch 2.1.0
+* Torch 2.0.1
 * BakLLaVA-1 model
 
 ## Available on RunPod