From 14ac2e985912872d16b9cdb28570d9a3f4b72635 Mon Sep 17 00:00:00 2001 From: Dhruv Nair Date: Tue, 31 Oct 2023 15:02:10 +0530 Subject: [PATCH] Update docker image for xformers (#5597) update docker image for xformers --- docker/diffusers-pytorch-xformers-cuda/Dockerfile | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docker/diffusers-pytorch-xformers-cuda/Dockerfile b/docker/diffusers-pytorch-xformers-cuda/Dockerfile index 95fe933798bc6..003f8e1165a1a 100644 --- a/docker/diffusers-pytorch-xformers-cuda/Dockerfile +++ b/docker/diffusers-pytorch-xformers-cuda/Dockerfile @@ -1,4 +1,4 @@ -FROM nvidia/cuda:11.7.1-cudnn8-runtime-ubuntu20.04 +FROM nvidia/cuda:12.1.0-runtime-ubuntu20.04 LABEL maintainer="Hugging Face" LABEL repository="diffusers" @@ -25,8 +25,8 @@ ENV PATH="/opt/venv/bin:$PATH" # pre-install the heavy dependencies (these can later be overridden by the deps from setup.py) RUN python3 -m pip install --no-cache-dir --upgrade pip && \ python3 -m pip install --no-cache-dir \ - torch==2.0.1 \ - torchvision==0.15.2 \ + torch \ + torchvision \ torchaudio \ invisible_watermark && \ python3 -m pip install --no-cache-dir \