diff --git a/Dockerfile b/Dockerfile index c93372a2..fa9a45bf 100644 --- a/Dockerfile +++ b/Dockerfile @@ -39,7 +39,9 @@ RUN cargo build --profile release-opt # Adapted from: https://github.com/pytorch/pytorch/blob/master/Dockerfile FROM nvidia/cuda:12.1.0-devel-ubuntu22.04 as pytorch-install +# NOTE: When updating PyTorch version, make sure to update server/pyproject.toml's `nvidia-nccl-cu12` requirement. Context: https://github.com/huggingface/text-generation-inference/pull/2099 ARG PYTORCH_VERSION=2.3.0 + ARG PYTHON_VERSION=3.10 # Keep in sync with `server/pyproject.toml ARG CUDA_VERSION=12.1