Browse Source

[Docker] Fix a couple of build issues (#3691)

pull/3824/head^2
Yanming W 2 years ago committed by GitHub
parent
commit
269150b6f4
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
  1. 4
      docker/Dockerfile
  2. 2
      op_builder/utils.py

4
docker/Dockerfile

@ -8,9 +8,13 @@ LABEL org.opencontainers.image.base.name = "docker.io/library/hpcaitech/cuda-con
# install torch
RUN conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.3 -c pytorch
# install ninja
RUN apt-get install -y --no-install-recommends ninja-build
# install apex
RUN git clone https://github.com/NVIDIA/apex && \
cd apex && \
git checkout 91fcaa && \
pip install packaging && \
pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" --global-option="--fast_layer_norm" ./

2
op_builder/utils.py

@ -110,7 +110,7 @@ def get_pytorch_version() -> List[int]:
torch_version = torch.__version__.split('+')[0]
TORCH_MAJOR = int(torch_version.split('.')[0])
TORCH_MINOR = int(torch_version.split('.')[1])
TORCH_PATCH = int(torch_version.split('.')[2])
TORCH_PATCH = int(torch_version.split('.')[2], 16)
return TORCH_MAJOR, TORCH_MINOR, TORCH_PATCH

Loading…
Cancel
Save