diff --git a/Dockerfile b/Dockerfile
index 461473a..b282831 100644
--- 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
index bec4163..9bec5b3 100644
--- 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