Skip to content

Commit 2ec8ae9

Browse files
authored
Use mamba & include cuda upgrades in build
1 parent 6732bc3 commit 2ec8ae9

File tree

1 file changed

+11
-3
lines changed

1 file changed

+11
-3
lines changed

packages/torch.Dockerfile

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,17 @@ ARG CUDA_MINOR_VERSION
1212
# TORCHVISION_VERSION is mandatory
1313
RUN test -n "$TORCHVISION_VERSION"
1414

15+
# Use mamba to speed up conda installs
16+
RUN conda install -c conda-forge mamba
17+
18+
# Install cudf/cuml so that cudatoolkit upgrades are included in the pytorch build
19+
RUN conda config --add channels nvidia && \
20+
conda config --add channels rapidsai
21+
RUN mamba install -y cudf cuml
22+
1523
# Build instructions: https://github.com/pytorch/pytorch#from-source
16-
RUN conda install astunparse numpy ninja pyyaml mkl mkl-include setuptools cmake cffi typing_extensions future six requests dataclasses
17-
RUN conda install -c pytorch magma-cuda${CUDA_MAJOR_VERSION}${CUDA_MINOR_VERSION}
24+
RUN mamba install astunparse numpy ninja pyyaml mkl mkl-include setuptools cmake cffi typing_extensions future six requests dataclasses
25+
RUN mamba install -c pytorch magma-cuda${CUDA_MAJOR_VERSION}${CUDA_MINOR_VERSION}
1826

1927
# By default, it uses the version from version.txt which includes the `a0` (alpha zero) suffix and part of the git hash.
2028
# This causes dependency conflicts like these: https://paste.googleplex.com/4786486378496000
@@ -46,7 +54,7 @@ RUN sudo apt-get update && \
4654
# ncurses.h is required for this install
4755
sudo apt-get install libncurses-dev && \
4856
# Fixing the build: https://github.com/pytorch/audio/issues/666#issuecomment-635928685
49-
conda install -c conda-forge ncurses && \
57+
mamba install -c conda-forge ncurses && \
5058
cd /usr/local/src && \
5159
git clone https://github.com/pytorch/audio && \
5260
cd audio && \

0 commit comments

Comments
 (0)