force cudatoolkit 11.2.2

Author: djherbis

Dockerfile.tmpl

@@ -90,7 +90,9 @@ RUN conda config --add channels nvidia && \
 # b/232247930: uninstall pyarrow to avoid double installation with the GPU specific version.
 RUN pip uninstall -y pyarrow && \
     conda install -c conda-forge mamba && \
-    mamba install -y cudf cuml cudatoolkit==11.2.2 && \
+    mamba install -y cudf cuml && \
+    # b/255757999: Force switch the cudatoolkit version back to the existing installation so it is compatible with pytorch
+    conda remove -y --force cudatoolkit=11.2.72 && mamba install --channel https://conda.anaconda.org/conda-forge/linux-64 cudatoolkit=11.2.2 && \
     /tmp/clean-layer.sh
 {{ end }}