fpreiss commited on
Commit
8056ecd
1 Parent(s): 738a057

add deepspeed-kernels dependency for deepspeed>=0.12.0 (#827)

Browse files
Files changed (1) hide show
  1. docker/Dockerfile-base +1 -1
docker/Dockerfile-base CHANGED
@@ -27,7 +27,7 @@ ENV PATH="/root/miniconda3/envs/py${PYTHON_VERSION}/bin:${PATH}"
27
  WORKDIR /workspace
28
 
29
  RUN python3 -m pip install --upgrade pip && pip3 install packaging && \
30
- python3 -m pip install --no-cache-dir -U torch==${PYTORCH_VERSION}+cu${CUDA} --extra-index-url https://download.pytorch.org/whl/cu$CUDA
31
 
32
  FROM base-builder AS deepspeed-builder
33
 
 
27
  WORKDIR /workspace
28
 
29
  RUN python3 -m pip install --upgrade pip && pip3 install packaging && \
30
+ python3 -m pip install --no-cache-dir -U torch==${PYTORCH_VERSION}+cu${CUDA} deepspeed-kernels --extra-index-url https://download.pytorch.org/whl/cu$CUDA
31
 
32
  FROM base-builder AS deepspeed-builder
33