diff --git a/setup.py b/setup.py
index dd1015d6..d3360850 100644
--- a/setup.py
+++ b/setup.py
@@ -34,6 +34,7 @@
 SUPPORTED_ARCHS = {"8.0", "8.6", "8.9", "9.0", "12.0"}
 
 # Compiler flags.
+TORCH_CUDA_ARCH_LIST = os.environ.get("TORCH_CUDA_ARCH_LIST", "")
 CXX_FLAGS = ["-g", "-O3", "-fopenmp", "-lgomp", "-std=c++17", "-DENABLE_BF16"]
 NVCC_FLAGS = [
     "-O3",
@@ -76,6 +77,9 @@ def get_nvcc_cuda_version(cuda_dir: str) -> Version:
         continue
     compute_capabilities.add(f"{major}.{minor}")
 
+if TORCH_CUDA_ARCH_LIST:
+    compute_capabilities = set(TORCH_CUDA_ARCH_LIST.split(","))
+
 nvcc_cuda_version = get_nvcc_cuda_version(CUDA_HOME)
 if not compute_capabilities:
     raise RuntimeError("No GPUs found. Please specify the target GPU architectures or build on a machine with GPUs.")