diff --git a/.circleci/unittest/linux/scripts/install.sh b/.circleci/unittest/linux/scripts/install.sh index 7e7461de8af..ace8583e47b 100755 --- a/.circleci/unittest/linux/scripts/install.sh +++ b/.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 diff --git a/.circleci/unittest/windows/scripts/install.sh b/.circleci/unittest/windows/scripts/install.sh index 17f32897165..cfdff3da6ba 100644 --- a/.circleci/unittest/windows/scripts/install.sh +++ b/.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}" diff --git a/packaging/build_conda.sh b/packaging/build_conda.sh index 4215c31016b..7c45aa3e6d9 100755 --- a/packaging/build_conda.sh +++ b/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 diff --git a/packaging/pkg_helpers.bash b/packaging/pkg_helpers.bash index 8cfb9e61b65..535564fbcc7 100644 --- a/packaging/pkg_helpers.bash +++ b/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"