From 6726a71bf6a0f2a52a2301fe3c04376f71340c40 Mon Sep 17 00:00:00 2001
From: Atinoda <61033436+Atinoda@users.noreply.github.com>
Date: Mon, 11 Mar 2024 21:39:11 +0000
Subject: [PATCH] Bump torch versions

Closes #44
---
 Dockerfile | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 97cbc80..1ff2447 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -37,7 +37,7 @@ RUN cp -ar /src /app
 # Base
 FROM app_base AS app_nvidia
 # Install pytorch for CUDA 12.1
-RUN pip3 install torch==2.1.* torchvision==0.16.* torchaudio==2.1.* \
+RUN pip3 install torch==2.2.1 torchvision==0.17.1 torchaudio==2.2.1 \
     --index-url https://download.pytorch.org/whl/cu121 
 # Install oobabooga/text-generation-webui
 RUN pip3 install -r /app/requirements.txt
@@ -53,7 +53,7 @@ RUN chmod +x /scripts/build_extensions.sh && \
 # Base
 FROM app_base AS app_rocm
 # Install pytorch for ROCM
-RUN pip3 install torch==2.1.* torchvision==0.16.* torchaudio==2.1.* \
+RUN pip3 install torch==2.2.1 torchvision==0.17.1 torchaudio==2.2.1 \
     --index-url https://download.pytorch.org/whl/rocm5.6
 # Install oobabooga/text-generation-webui
 RUN pip3 install -r /app/requirements_amd.txt
@@ -89,7 +89,7 @@ RUN chmod +x /scripts/build_extensions.sh && \
 # Base
 FROM app_base AS app_cpu
 # Install pytorch for CPU
-RUN pip3 install torch==2.1.* torchvision==0.16.* torchaudio==2.1.* \
+RUN pip3 install torch==2.2.1 torchvision==0.17.1 torchaudio==2.2.1 \
     --index-url https://download.pytorch.org/whl/cpu
 # Install oobabooga/text-generation-webui
 RUN pip3 install -r /app/requirements_cpu_only.txt
@@ -217,4 +217,4 @@ COPY --from=app_cpu_x $VIRTUAL_ENV $VIRTUAL_ENV
 # Variant parameters
 RUN echo "CPU Extended" > /variant.txt
 ENV EXTRA_LAUNCH_ARGS=""
-CMD ["python3", "/app/server.py"]
\ No newline at end of file
+CMD ["python3", "/app/server.py"]