diff --git a/docker/Dockerfile b/docker/Dockerfile index 8d58449..7f56e9c 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -97,6 +97,8 @@ RUN python${PYTHON_VER} -m venv ${VENV} && \ mkdir -p /root/.pip && \ echo "[global]" > /root/.pip/pip.conf && \ echo "index-url = ${PIP_INDEX}" >> /root/.pip/pip.conf && \ + echo "[install]" >> /root/.pip/pip.conf && \ + echo "constraint = ${VENV}/torch-constraint.txt" >> /root/.pip/pip.conf && \ ${VENV}/bin/pip install -U pip setuptools wheel && \ # PyTorch ROCm wheels 在海外 CDN,通过代理下载 export http_proxy=${GIT_PROXY} https_proxy=${GIT_PROXY} && \ @@ -105,6 +107,8 @@ RUN python${PYTHON_VER} -m venv ${VENV} && \ torchvision \ pytorch-triton-rocm \ --index-url ${TORCH_INDEX} && \ + # 锁定 torch 版本,防止后续 pip install 解析到 CUDA 版 + echo "torch==2.11.0+rocm7.2" > ${VENV}/torch-constraint.txt && \ ${VENV}/bin/python -c "import torch; print('PyTorch:', torch.__version__); assert torch.version.hip is not None" # ===========================================================================