Add CUDA 11.6 version to manywheel script

Make sure TORCH_CUDA_ARCH_LIST and EXTRA_CAFFE2_CMAKE_FLAGS are set for CUDA 11.6
pytorch · Apr 8, 2022 · 7d5e98f · 7d5e98f
1 parent 96a30f6
commit 7d5e98f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/manywheel/build_cuda.sh b/manywheel/build_cuda.sh
@@ -58,7 +58,7 @@ cuda_version_nodot=$(echo $CUDA_VERSION | tr -d '.')
 
 TORCH_CUDA_ARCH_LIST="3.7;5.0;6.0;7.0"
 case ${CUDA_VERSION} in
-    11.[35])
+    11.[356])
         TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST};7.5;8.0;8.6"
         EXTRA_CAFFE2_CMAKE_FLAGS+=("-DATEN_NO_TEST=ON")
         ;;