update

pyg-team · Oct 5, 2023 · 8b30c44 · 8b30c44
1 parent 915e8f7
commit 8b30c44
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/setup.py b/setup.py
@@ -64,6 +64,10 @@ def build_extension(self, ext):
             f'-DCMAKE_PREFIX_PATH={torch.utils.cmake_prefix_path}',
         ]
 
+        arch_list = os.environ.get('TORCH_CUDA_ARCH_LIST')
+        if WITH_CUDA and arch_list is not None:
+            cmake_args.append(f'-DCMAKE_CUDA_ARCHITECTURES={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}')