diff --git a/packaging/pkg_helpers.bash b/packaging/pkg_helpers.bash index 897cc103345..d670f789548 100644 --- a/packaging/pkg_helpers.bash +++ b/packaging/pkg_helpers.bash @@ -52,7 +52,6 @@ setup_cuda() { else export CUDA_HOME=/usr/local/cuda-11.3/ fi - export FORCE_CUDA=1 export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5;8.0;8.6" ;; cu112) @@ -61,7 +60,6 @@ setup_cuda() { else export CUDA_HOME=/usr/local/cuda-11.2/ fi - export FORCE_CUDA=1 export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5;8.0;8.6" ;; cu111) @@ -70,7 +68,6 @@ setup_cuda() { else export CUDA_HOME=/usr/local/cuda-11.1/ fi - export FORCE_CUDA=1 export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5;8.0;8.6" ;; cu110) @@ -79,7 +76,6 @@ setup_cuda() { else export CUDA_HOME=/usr/local/cuda-11.0/ fi - export FORCE_CUDA=1 export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5;8.0" ;; cu102) @@ -88,7 +84,6 @@ setup_cuda() { else export CUDA_HOME=/usr/local/cuda-10.2/ fi - export FORCE_CUDA=1 export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5" ;; cu101) @@ -97,7 +92,6 @@ setup_cuda() { else export CUDA_HOME=/usr/local/cuda-10.1/ fi - export FORCE_CUDA=1 export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5" ;; cu100) @@ -106,7 +100,6 @@ setup_cuda() { else export CUDA_HOME=/usr/local/cuda-10.0/ fi - export FORCE_CUDA=1 export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5" ;; cu92) @@ -115,7 +108,6 @@ setup_cuda() { else export CUDA_HOME=/usr/local/cuda-9.2/ fi - export FORCE_CUDA=1 export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0" ;; cpu) @@ -128,6 +120,11 @@ setup_cuda() { exit 1 ;; esac + if [[ -n "$CUDA_HOME" ]]; then + # Adds nvcc binary to the search path so that CMake's `find_package(CUDA)` will pick the right one + export PATH="$CUDA_HOME/bin:$PATH" + export FORCE_CUDA=1 + fi } # Populate build version if necessary, and add version suffix