From 8f842f7291c56086de4c67b3cb7234537cf72d1d Mon Sep 17 00:00:00 2001 From: philschmid Date: Sun, 18 Dec 2022 09:58:57 +0000 Subject: [PATCH] pin cuda and torchvision --- dockerfiles/starlette/pytorch/Dockerfile.gpu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dockerfiles/starlette/pytorch/Dockerfile.gpu b/dockerfiles/starlette/pytorch/Dockerfile.gpu index 89be75d1..3cea926c 100644 --- a/dockerfiles/starlette/pytorch/Dockerfile.gpu +++ b/dockerfiles/starlette/pytorch/Dockerfile.gpu @@ -5,7 +5,7 @@ COPY starlette_requirements.txt /tmp/requirements.txt RUN pip install --no-cache-dir -r /tmp/requirements.txt && rm /tmp/requirements.txt # Think about a better solution -> base contaienr has pt 1.13. thats why need below 0.14 -RUN pip install --no-cache-dir sentence_transformers torchvision~="0.14.0" diffusers=="0.9.0" accelerate=="0.14.0" +RUN pip install --no-cache-dir sentence_transformers torchvision=="0.14.0" diffusers=="0.9.0" accelerate=="0.14.0" --extra-index-url https://download.pytorch.org/whl/cu116 # Add upgrade due to issue in base container upgrade https://github.com/mamba-org/mamba/issues/2170 RUN pip install transformers==4.25.1 --no-cache-dir --upgrade