1
0
Эх сурвалжийг харах

[LayerNorm] Compile for both sm70 and sm80

Tri Dao 2 жил өмнө
parent
commit
39ed597b28

+ 2 - 2
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")