Fix unhandled cuda error for Conda build with CUDA 12.6 (#2025)

related pytorch/pytorch#138440
fix error for https://github.com/pytorch/pytorch/pull/138899/files

conda/build_pytorch.sh

@@ -268,7 +268,10 @@ else
     . ./switch_cuda_version.sh "$desired_cuda"
     # TODO, simplify after anaconda fixes their cudatoolkit versioning inconsistency.
     # see: https://github.com/conda-forge/conda-forge.github.io/issues/687#issuecomment-460086164
-    if [[ "$desired_cuda" == "12.4" ]]; then
+    if [[ "$desired_cuda" == "12.6" ]]; then
+        export CONDA_CUDATOOLKIT_CONSTRAINT="    - pytorch-cuda >=12.6,<12.7 # [not osx]"
+        export MAGMA_PACKAGE="    - magma-cuda126 # [not osx and not win]"
+    elif [[ "$desired_cuda" == "12.4" ]]; then
         export CONDA_CUDATOOLKIT_CONSTRAINT="    - pytorch-cuda >=12.4,<12.5 # [not osx]"
         export MAGMA_PACKAGE="    - magma-cuda124 # [not osx and not win]"
     elif [[ "$desired_cuda" == "12.1" ]]; then

conda/pytorch-nightly/build.sh

@@ -60,7 +60,7 @@ if [[ -n "$build_with_cuda" ]]; then
         TORCH_CUDA_ARCH_LIST="$TORCH_CUDA_ARCH_LIST;3.7+PTX;9.0"
         #for cuda 11.8 include all dynamic loading libraries
         DEPS_LIST=(/usr/local/cuda/lib64/libcudnn*.so.9 /usr/local/cuda-11.8/extras/CUPTI/lib64/libcupti.so.11.8 /usr/local/cuda/lib64/libcusparseLt.so.0)
-    elif [[ $CUDA_VERSION == 12.1* || $CUDA_VERSION == 12.4* ]]; then
+    elif [[ $CUDA_VERSION == 12.1* || $CUDA_VERSION == 12.4* || $CUDA_VERSION == 12.6*]]; then
         # cuda 12 does not support sm_3x
         TORCH_CUDA_ARCH_LIST="$TORCH_CUDA_ARCH_LIST;9.0"
         # for cuda 12.1 (12.4) we use cudnn 9.1 and include all dynamic loading libraries