You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
# the remaining pip commands: https://www.anaconda.com/using-pip-in-a-conda-environment/
56
58
# However, because this image is based on the CPU image, this isn't possible but better
57
59
# to put them at the top of this file to minize conflicts.
58
-
RUN conda install cudf=21.06 cuml=21.06 cudatoolkit=$CUDA_VERSION && \
60
+
RUN conda install cudf=21.08 cuml=21.08 cudatoolkit=$CUDA_VERSION && \
59
61
/tmp/clean-layer.sh
60
62
61
63
# Install Pytorch and torchvision with GPU support.
62
-
# Note: torchtext and torchaudio do not require a separate GPU package.
63
-
RUN pip install torch==1.7.1+cu$CUDA_MAJOR_VERSION$CUDA_MINOR_VERSION torchvision==0.8.2+cu$CUDA_MAJOR_VERSION$CUDA_MINOR_VERSION -f https://download.pytorch.org/whl/torch_stable.html && \
64
+
# Note: torchtext and torchaudio do not require a separate package.
65
+
# Replace `cu111` by `cu$CUDA_MAJOR_VERSION$CUDA_MINOR_VERSION` once build for CUDA 11.2 is released.
66
+
# Introduced in CUDA 11.1, CUDA Enhanced Compatibility leverages semantic versioning across components in the CUDA Toolkit, an application can be built for one CUDA minor release (such as 11.1) and work across all future minor releases within the major family (such as 11.x).
# Install JAX (Keep JAX version in sync with CPU image)
82
-
RUN pip install jax==0.2.16 jaxlib==0.1.68+cuda$CUDA_MAJOR_VERSION$CUDA_MINOR_VERSION -f https://storage.googleapis.com/jax-releases/jax_releases.html && \
87
+
# TODO(b/181966788) Replace `cuda111` with `cuda$CUDA_MAJOR_VERSION$CUDA_MINOR_VERSION` once new version is out.
88
+
RUN pip install jax==0.2.16 jaxlib==0.1.68+cuda111 -f https://storage.googleapis.com/jax-releases/jax_releases.html && \
83
89
/tmp/clean-layer.sh
84
90
85
91
# Reinstall packages with a separate version for GPU support.
@@ -90,9 +96,8 @@ RUN pip uninstall -y mxnet && \
90
96
# Install GPU-only packages
91
97
RUN pip install pycuda && \
92
98
pip install pynvrtc && \
93
-
# b/190622765 latest version is causing issue. nnabla fixed it in https://github.com/sony/nnabla/issues/892, waiting for new release before we can remove this pin.
0 commit comments