Allow cmake vars in docker build (#100867)

Fixes #100866

Pull Request resolved: https://github.com/pytorch/pytorch/pull/100867
Approved by: https://github.com/malfet
This commit is contained in:
Kristoffer Carlson
2023-05-10 17:44:26 +00:00
committed by PyTorch MergeBot
parent 058d740f59
commit e762cce61f
3 changed files with 13 additions and 1 deletions

View File

@ -50,10 +50,12 @@ COPY . .
RUN git submodule update --init --recursive
FROM conda as build
ARG CMAKE_VARS
WORKDIR /opt/pytorch
COPY --from=conda /opt/conda /opt/conda
COPY --from=submodule-update /opt/pytorch /opt/pytorch
RUN --mount=type=cache,target=/opt/ccache \
export eval ${CMAKE_VARS} && \
TORCH_CUDA_ARCH_LIST="3.5 5.2 6.0 6.1 7.0+PTX 8.0" TORCH_NVCC_FLAGS="-Xfatbin -compress-all" \
CMAKE_PREFIX_PATH="$(dirname $(which conda))/../" \
python setup.py install