Skip to content

GitLab

"docs/git@developer.sourcefind.cn:OpenDAS/mmcv.git" did not exist on "d510b8b174645303fe2ad3cc43bbcde3c752cec5"
Unverified Commit e1aa1fc1 authored by Masaki Kozuki's avatar Masaki Kozuki Committed by GitHub
Browse files

8.6 requires CUDA 11.1 (#1289)

parent a786ca0c
Hide whitespace changes
Inline Side-by-side
Showing with 4 additions and 2 deletions
setup.py
View file @ e1aa1fc1
...@@ -72,9 +72,11 @@ if not torch.cuda.is_available(): ...@@ -72,9 +72,11 @@ if not torch.cuda.is_available():
'export TORCH_CUDA_ARCH_LIST="compute capability" before running setup.py.\n', 'export TORCH_CUDA_ARCH_LIST="compute capability" before running setup.py.\n',
) )
if os.environ.get("TORCH_CUDA_ARCH_LIST", None) is None: if os.environ.get("TORCH_CUDA_ARCH_LIST", None) is None:
_, bare_metal_major, _ = get_cuda_bare_metal_version(CUDA_HOME) _, bare_metal_major, bare_metal_minor = get_cuda_bare_metal_version(CUDA_HOME)
if int(bare_metal_major) == 11: if int(bare_metal_major) == 11:
os.environ["TORCH_CUDA_ARCH_LIST"] = "6.0;6.1;6.2;7.0;7.5;8.0;8.6" os.environ["TORCH_CUDA_ARCH_LIST"] = "6.0;6.1;6.2;7.0;7.5;8.0"
if int(bare_metal_minor) > 0:
os.environ["TORCH_CUDA_ARCH_LIST"] = "6.0;6.1;6.2;7.0;7.5;8.0;8.6"
else: else:
os.environ["TORCH_CUDA_ARCH_LIST"] = "6.0;6.1;6.2;7.0;7.5" os.environ["TORCH_CUDA_ARCH_LIST"] = "6.0;6.1;6.2;7.0;7.5"
......
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