From 645cfce64cd5d3ecf85fbbef1b6e512c90bffadb Mon Sep 17 00:00:00 2001 From: Pengfei Xuan Date: Mon, 6 May 2024 17:17:25 -0400 Subject: [PATCH] Sync TORCH_CUDA_ARCH_LIST from CMAKE_CUDA_ARCHITECTURES --- CMakeLists.txt | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/CMakeLists.txt b/CMakeLists.txt index 24026c6..3841b79 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -107,6 +107,11 @@ if(GPU_RUNTIME STREQUAL "CUDA") endforeach() message("Set CUDA flags: " ${CMAKE_CUDA_FLAGS}) endif() + # Set torch cuda architecture list + set(TORCH_CUDA_ARCH_LIST ${CMAKE_CUDA_ARCHITECTURES}) + list(TRANSFORM TORCH_CUDA_ARCH_LIST REPLACE "([0-9])([0-9])" "\\1.\\2") + string(REPLACE ";" " " TORCH_CUDA_ARCH_LIST "${TORCH_CUDA_ARCH_LIST}") + message(STATUS "** Updated TORCH_CUDA_ARCH_LIST to ${TORCH_CUDA_ARCH_LIST} **") endif() elseif(GPU_RUNTIME STREQUAL "HIP") set(USE_HIP ON CACHE BOOL "Use HIP for GPU acceleration")