Replacing cudatoolkit by cuda for 11.6 (#5996)

Author: atalman

.circleci/unittest/linux/scripts/install.sh

@@ -21,7 +21,12 @@ else
     fi
     echo "Using CUDA $CUDA_VERSION as determined by CU_VERSION"
     version="$(python -c "print('.'.join(\"${CUDA_VERSION}\".split('.')[:2]))")"
-    cudatoolkit="nvidia::cudatoolkit=${version}"
+
+    cuda_toolkit_pckg="cudatoolkit"
+    if [[ "$CU_VERSION" == cu116 ]]; then
+        cuda_toolkit_pckg="cuda"
+    fi
+    cudatoolkit="nvidia::${cuda_toolkit_pckg}=${version}"
 fi
 
 case "$(uname -s)" in

.circleci/unittest/windows/scripts/install.sh

@@ -22,9 +22,15 @@ else
     elif [[ ${#CU_VERSION} -eq 5 ]]; then
         CUDA_VERSION="${CU_VERSION:2:2}.${CU_VERSION:4:1}"
     fi
+
+    cuda_toolkit_pckg="cudatoolkit"
+    if [[ "$CU_VERSION" == cu116 ]]; then
+        cuda_toolkit_pckg="cuda"
+    fi
+
     echo "Using CUDA $CUDA_VERSION as determined by CU_VERSION"
     version="$(python -c "print('.'.join(\"${CUDA_VERSION}\".split('.')[:2]))")"
-    cudatoolkit="cudatoolkit=${version}"
+    cudatoolkit="${cuda_toolkit_pckg}=${version}"
 fi
 
 printf "Installing PyTorch with %s\n" "${cudatoolkit}"

packaging/build_conda.sh

@@ -11,11 +11,6 @@ setup_conda_pytorch_constraint
 setup_conda_cudatoolkit_constraint
 setup_visual_studio_constraint
 setup_junit_results_folder
-
-# nvidia channel included for cudatoolkit >= 11 however for 11.5 and 11.6 we use conda-forge
 export CUDATOOLKIT_CHANNEL="nvidia"
-if [[ "$CU_VERSION" == cu116 ]]; then
-    export CUDATOOLKIT_CHANNEL="conda-forge"
-fi
 
 conda build -c $CUDATOOLKIT_CHANNEL -c defaults $CONDA_CHANNEL_FLAGS --no-anaconda-upload --python "$PYTHON_VERSION" packaging/torchvision

packaging/pkg_helpers.bash

@@ -257,7 +257,7 @@ setup_conda_cudatoolkit_constraint() {
   else
     case "$CU_VERSION" in
       cu116)
-        export CONDA_CUDATOOLKIT_CONSTRAINT="- cudatoolkit >=11.6,<11.7 # [not osx]"
+        export CONDA_CUDATOOLKIT_CONSTRAINT="- cuda >=11.6,<11.7 # [not osx]"
         ;;
       cu113)
         export CONDA_CUDATOOLKIT_CONSTRAINT="- cudatoolkit >=11.3,<11.4 # [not osx]"
@@ -286,7 +286,7 @@ setup_conda_cudatoolkit_plain_constraint() {
   else
     case "$CU_VERSION" in
       cu116)
-        export CONDA_CUDATOOLKIT_CONSTRAINT="cudatoolkit=11.6"
+        export CONDA_CUDATOOLKIT_CONSTRAINT="cuda=11.6"
         ;;
       cu113)
         export CONDA_CUDATOOLKIT_CONSTRAINT="cudatoolkit=11.3"