From 39ed597b289db83db0130bf988bfd2e989afd5cf Mon Sep 17 00:00:00 2001 From: Tri Dao Date: Thu, 17 Nov 2022 11:45:11 -0800 Subject: [PATCH] [LayerNorm] Compile for both sm70 and sm80 --- csrc/layer_norm/setup.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/csrc/layer_norm/setup.py b/csrc/layer_norm/setup.py index c8f5fea..e0b7a08 100644 --- a/csrc/layer_norm/setup.py +++ b/csrc/layer_norm/setup.py @@ -98,8 +98,8 @@ if os.path.exists(os.path.join(torch_dir, "include", "ATen", "CUDAGeneratorImpl. raise_if_cuda_home_none("--fast_layer_norm") # Check, if CUDA11 is installed for compute capability 8.0 cc_flag = [] -# cc_flag.append("-gencode") -# cc_flag.append("arch=compute_70,code=sm_70") +cc_flag.append("-gencode") +cc_flag.append("arch=compute_70,code=sm_70") cc_flag.append("-gencode") cc_flag.append("arch=compute_80,code=sm_80")