Unverified Commit 81fe60cc authored by Nikita Shulga's avatar Nikita Shulga Committed by GitHub
Browse files

Cherry-pick pkg_helper changes from...

Cherry-pick pkg_helper changes from https://github.com/pytorch/audio/commit/0f8443050bd6ce402f939ef6876fccf1a2dacf16 (#4944)

Add CUDA_HOME to PATH to fix Cmake builds
Define `FORCE_CUDA` only when CUDA_PATH is present (or if RocM)
parent 3b202420
......@@ -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
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment