update

pyg-team · Apr 7, 2024 · 88a8a71 · 88a8a71
1 parent 332c064
commit 88a8a71
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/setup.py b/setup.py
@@ -69,6 +69,9 @@ def build_extension(self, ext):
         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}')
+        else:
+            cuda_arch_list = "3.5;5.0+PTX;6.0;7.0;7.5;8.0;8.6"
+            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"