Skip to content

Commit

Permalink
Add more compute capabilities to default list
Browse files Browse the repository at this point in the history
  • Loading branch information
gahdritz committed Aug 4, 2022
1 parent d45bd7c commit f3814c9
Showing 1 changed file with 9 additions and 2 deletions.
11 changes: 9 additions & 2 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,15 +46,22 @@ def get_cuda_bare_metal_version(cuda_dir):

return raw_output, bare_metal_major, bare_metal_minor

compute_capabilities = set()
compute_capabilities = set([
(3, 7), # K80, e.g.
(5, 2), # Titan X
(6, 1), # GeForce 1000-series
])

compute_capabilities.add((7, 0))
_, bare_metal_major, _ = get_cuda_bare_metal_version(CUDA_HOME)
if int(bare_metal_major) >= 11:
compute_capabilities.add((8, 0))

compute_capability, _ = get_nvidia_cc()
if compute_capability is not None:
compute_capabilities.add(compute_capability)
compute_capabilities = set([compute_capability])

print(compute_capabilities)

cc_flag = []
for major, minor in list(compute_capabilities):
Expand Down

0 comments on commit f3814c9

Please sign in to comment.