Update setup.py (#68)

revert back device setting
microsoft · Oct 16, 2024 · 994f9e4 · 994f9e4
1 parent b9d6a2c
commit 994f9e4
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/setup.py b/setup.py
@@ -22,7 +22,7 @@ def get_version():
 
 
 def build_cuda_extensions():
-    compute_capabilities = [80]
+    compute_capabilities = [70, 75, 80, 86, 89, 90]
     arch_flags = []
     TORCH_CUDA_ARCH_LIST = os.getenv("TORCH_CUDA_ARCH_LIST", None)
     if TORCH_CUDA_ARCH_LIST is None: