update

pyg-team · Apr 7, 2024 · 332c064 · 332c064
1 parent b95e7d0
commit 332c064
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/.github/workflows/cuda/Windows.sh b/.github/workflows/cuda/Windows.sh
@@ -72,5 +72,4 @@ cp -r /tmp/NvToolsExt/include/* "/c/Program Files/NVIDIA Corporation/NvToolsExt/
 cp -r /tmp/NvToolsExt/lib/x64/* "/c/Program Files/NVIDIA Corporation/NvToolsExt/lib/x64"
 export NVTOOLSEXT_PATH="/c/Program Files/NVIDIA Corporation/NvToolsExt/bin/x64"
 
-export CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5;8.0;8.6"
 export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5;8.0;8.6"
diff --git a/setup.py b/setup.py
@@ -66,6 +66,10 @@ def build_extension(self, ext):
             f'-DCMAKE_PREFIX_PATH={torch.utils.cmake_prefix_path}',
         ]
 
+        cuda_arch_list = os.getenv('TORCH_CUDA_ARCH_LIST')
+        if WITH_CUDA and cuda_arch_list is not None:
+            cmake_args.append(f'-DCMAKE_CUDA_ARCHITECTURES={cuda_arch_list}')
+
         if CMakeBuild.check_env_flag('USE_MKL_BLAS'):
             include_dir = f"{sysconfig.get_path('data')}{os.sep}include"
             cmake_args.append(f'-DBLAS_INCLUDE_DIR={include_dir}')