Update setup.py

detect device and set `--generate-code` automatically
databricks · Jul 26, 2023 · 7c9a26d · 7c9a26d
1 parent 0032d97
commit 7c9a26d
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/setup.py b/setup.py
@@ -1,7 +1,10 @@
 from setuptools import setup, find_packages
+from torch import cuda
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension
 
 
+_dc = cuda.get_device_capability()
+_dc = f"{_dc[0]}{_dc[1]}"
 ext_modules = [
     CUDAExtension(
         "megablocks_ops",
@@ -12,7 +15,7 @@
             "nvcc": [
                 "--ptxas-options=-v",
                 "--optimize=2",
-                "--generate-code=arch=compute_80,code=sm_80"
+                f"--generate-code=arch=compute_{_dc},code=sm_{_dc}"
             ]
         })
 ]