Allow to pass architectures like 90a, without being overriden

Signed-off-by: aurianer <[email protected]>
NVIDIA · Sep 11, 2024 · 1d74955 · 1d74955
1 parent 40dda92
commit 1d74955
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/build_tools/pytorch.py b/build_tools/pytorch.py
@@ -80,10 +80,10 @@ def setup_pytorch_extension(
             )
         )
 
-        if "80" in cuda_architectures:
-            nvcc_flags.extend(["-gencode", "arch=compute_80,code=sm_80"])
-        if "90" in cuda_architectures:
-            nvcc_flags.extend(["-gencode", "arch=compute_90,code=sm_90"])
+        for arch in cuda_architectures:
+            if arch_nb != "70":
+                arch_nb = arch.replace("sm_", "")
+                nvcc_flags.extend(["-gencode", f"arch=compute_{arch_nb},code={arch}"])
 
     # Libraries
     library_dirs = []