Search code examples
dockerpytorch

Installing pytorch is not the version I want

# Use Nvidia CUDA base image
FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04 as base

# Prevents prompts from packages asking for user input during installation
ENV DEBIAN_FRONTEND=noninteractive
# Prefer binary wheels over source distributions for faster pip installations
ENV PIP_PREFER_BINARY=1
# Ensures output from python is printed immediately to the terminal without buffering
ENV PYTHONUNBUFFERED=1 

# Install Python, git and other necessary tools
RUN apt-get update && apt-get install -y \
    python3.10 \
    python3-pip \
    git \
    wget \
    libgl-dev \
    python3-opencv \
    libglib2.0-0 \
    ffmpeg

# Clean up to reduce image size
RUN apt-get autoremove -y && apt-get clean -y && rm -rf /var/lib/apt/lists/*


# Install ComfyUI dependencies
RUN pip3 install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121 -i https://pypi.tuna.tsinghua.edu.cn/simple \
    && pip3 install --no-cache-dir xformers==0.0.21 -i https://pypi.tuna.tsinghua.edu.cn/simple

# Install runpod
RUN pip3 install runpod requests -i https://pypi.tuna.tsinghua.edu.cn/simple

# Go back to the root
WORKDIR /

# Add the start and the handler
ADD src/start.sh src/rp_handler.py test_input.json ./
RUN chmod +x /start.sh

# Start the container
CMD /start.sh

During the installation, I get an error message, see below enter image description here

After installing, I started the container and used torch.__version__ to check the version number and got 2.0.1

I should get pytorch version number 2.1.2

Solution

  • According to documentation, have you tried modifying this lines :

    RUN pip3 install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121 -i https://pypi.tuna.tsinghua.edu.cn/simple \
&& pip3 install --no-cache-dir xformers==0.0.21 -i https://pypi.tuna.tsinghua.edu.cn/simple

    to :

    RUN pip3 install torch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 --index-url https://download.pytorch.org/whl/test/cu118
&& pip3 install --no-cache-dir xformers==0.0.21 -i https://pypi.tuna.tsinghua.edu.cn/simple