Update builder.py (#642)

64461da4 · Xingjian Shi · GitHub · 44bd538b · 64461da4
Unverified Commit 64461da4 authored Jan 06, 2021 by Xingjian Shi Committed by GitHub Jan 06, 2021
Show whitespace changes
Inline Side-by-side

Showing with 5 additions and 1 deletion

op_builder/builder.py op_builder/builder.py +5 -1

No files found.
--- a/op_builder/builder.py
+++ b/op_builder/builder.py
@@ -36,6 +36,10 @@ def get_default_compute_capatabilities():
    import torch.utils.cpp_extension
    if torch.utils.cpp_extension.CUDA_HOME is not None and installed_cuda_version(
    )[0] >= 11:
+        if installed_cuda_version()[0] == 11 and installed_cuda_version()[1] == 0:
+            # Special treatment of CUDA 11.0 because compute_86 is not supported.
+            compute_caps += ";8.0"
+        else:
            compute_caps += ";8.0;8.6"
    return compute_caps