Skip to content
This repository was archived by the owner on Aug 15, 2025. It is now read-only.

manywheel: Fix TORCH_CUDA_ARCH_LIST not getting populated #609

Merged
merged 1 commit into from
Dec 10, 2020
Merged

manywheel: Fix TORCH_CUDA_ARCH_LIST not getting populated #609

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 8 additions & 8 deletions manywheel/build.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,22 +50,22 @@ fi

cuda_version_nodot=$(echo $CUDA_VERSION | tr -d '.')

torch_cuda_arch_list="3.7;5.0;6.0;7.0"
TORCH_CUDA_ARCH_LIST="3.7;5.0;6.0;7.0"
case ${CUDA_VERSION} in
11.1)
torch_cuda_arch_list="5.0;7.0;8.0;8.6" # removing some to prevent bloated binary size
TORCH_CUDA_ARCH_LIST="5.0;7.0;8.0;8.6" # removing some to prevent bloated binary size
EXTRA_CAFFE2_CMAKE_FLAGS+=("-DATEN_NO_TEST=ON")
;;
11.0)
torch_cuda_arch_list="$TORCH_CUDA_ARCH_LIST;7.5;8.0"
TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST};7.5;8.0"
EXTRA_CAFFE2_CMAKE_FLAGS+=("-DATEN_NO_TEST=ON")
;;
10.*)
torch_cuda_arch_list="$TORCH_CUDA_ARCH_LIST;7.5"
TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST};7.5"
EXTRA_CAFFE2_CMAKE_FLAGS+=("-DATEN_NO_TEST=ON")
;;
9.*)
torch_cuda_arch_list="${TORCH_CUDA_ARCH_LIST}"
TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}"
;;
*)
echo "unknown cuda version $CUDA_VERSION"
Expand All @@ -74,7 +74,7 @@ case ${CUDA_VERSION} in
esac

if [[ -n "$OVERRIDE_TORCH_CUDA_ARCH_LIST" ]]; then
torch_cuda_arch_list="$OVERRIDE_TORCH_CUDA_ARCH_LIST"
TORCH_CUDA_ARCH_LIST="$OVERRIDE_TORCH_CUDA_ARCH_LIST"

# Prune CUDA again with new arch list. Unfortunately, we need to re-install CUDA to prune it again
override_gencode=""
Expand All @@ -87,8 +87,8 @@ if [[ -n "$OVERRIDE_TORCH_CUDA_ARCH_LIST" ]]; then
bash "$(dirname "$SCRIPTPATH")"/common/install_cuda.sh "${CUDA_VERSION}"
fi

export TORCH_CUDA_ARCH_LIST=$torch_cuda_arch_list
echo "$TORCH_CUDA_ARCH_LIST"
export TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST}
echo "${TORCH_CUDA_ARCH_LIST}"

# Package directories
WHEELHOUSE_DIR="wheelhouse$cuda_version_nodot"
Expand Down