"examples/git@developer.sourcefind.cn:OpenDAS/colossalai.git" did not exist on "24cbee0ebe82efee10c8ac90379a15d8b8b1b016"
Unverified Commit cfa60708 authored by Mashiro's avatar Mashiro Committed by GitHub
Browse files

[Fix] Fix compile error (#4357)

parent eb952ea8
......@@ -197,11 +197,12 @@ def get_cuda_cc_flag() -> List[str]:
import torch
cc_flag = []
max_arch = ''.join(str(i) for i in torch.cuda.get_device_capability())
for arch in torch.cuda.get_arch_list():
res = re.search(r'sm_(\d+)', arch)
if res:
arch_cap = res[1]
if int(arch_cap) >= 60:
if int(arch_cap) >= 60 and int(arch_cap) <= int(max_arch):
cc_flag.extend(['-gencode', f'arch=compute_{arch_cap},code={arch}'])
return cc_flag
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment