Installing PyTorch with Python version 3.11 in a Docker container

Question

I see on the official PyTorch page that PyTorch supports Python versions 3.8 to 3.11.

When I actually try to install PyTorch + CUDA in a Python 3.11 Docker image, it seems unable to find CUDA drivers, e.g.

FROM python:3.11.4
RUN --mount=type=cache,id=pip-build,target=/root/.cache/pip \
    pip install torch torchaudio
ENV PATH="/usr/local/nvidia/bin:${PATH}" \
    NVIDIA_VISIBLE_DEVICES=all \
    NVIDIA_DRIVER_CAPABILITIES=all

Then, inside the container, I see that torch.version.cuda is None

Compare this to

FROM pytorch/pytorch
RUN --mount=type=cache,id=pip-build,target=/root/.cache/pip \
    pip install torchaudio
ENV PATH="/usr/local/nvidia/bin:${PATH}" \
    NVIDIA_VISIBLE_DEVICES=all \
    NVIDIA_DRIVER_CAPABILITIES=all

Inside the container I see that torch.version.cuda is 12.1

PyTorch claims they're compatible with Python 3.11, has anybody actually been able to use PyTorch+CUDA with Python 3.11?

Tried running Docker images with Python 3.11.4

Tried running the Conda docker image and installing pytorch, but kept getting errors that the images couldn't be found

Answer 1

Update - think I solved it myself. Kind of dumb, but I just upgrade the Python version in my Dockerfile. I'm sure this results in something breaking, but I haven't run into it yet.

FROM pytorch/pytorch

RUN conda update -n base -c defaults conda && \
    conda install -y python=3.11 && \
    conda update --all --yes

CMD ["python", "--version"]