Add cuda archs

openmm · Jun 27, 2023 · 7571e13 · 7571e13
1 parent 5779ad2
commit 7571e13
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/setup.py b/setup.py
@@ -153,7 +153,7 @@ def build_extension(self, ext: CMakeExtension) -> None:
     else:
         raise RuntimeError("Unsupported CUDA version")
     CMAKE_CUDA_ARCHS = ";".join([str(arch) for arch in ARCHES] + [f"{LATEST_ARCH}-real", f"{LATEST_ARCH}-virtual"])
-    extra_args["CMAKE_CUDA_ARCHITECTURES"] = "50"
+    extra_args["CMAKE_CUDA_ARCHITECTURES"] = CMAKE_CUDA_ARCHS
 else:
     extra_args["ENABLE_CUDA"] = "OFF"