Skip to content

GitLab

Unverified Commit 247349f1 authored by mcarilli's avatar mcarilli Committed by GitHub
Browse files

Update setup.py for #23

parent d74fda26
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 2 deletions
setup.py
View file @ 247349f1
...@@ -82,7 +82,7 @@ if TORCH_MAJOR == 0 and TORCH_MINOR == 4: ...@@ -82,7 +82,7 @@ if TORCH_MAJOR == 0 and TORCH_MINOR == 4:
def get_cuda_version(): def get_cuda_version():
NVCC = find(CUDA_HOME+os.sep+"bin", NVCC = find(CUDA_HOME+os.sep+"bin",
re.compile('nvcc$').search) re.compile('nvcc$|nvcc.exe').search)
print("Found NVCC = ", NVCC) print("Found NVCC = ", NVCC)
# Parse output of nvcc to get cuda major version # Parse output of nvcc to get cuda major version
...@@ -104,7 +104,8 @@ if CUDA_HOME is not None: ...@@ -104,7 +104,8 @@ if CUDA_HOME is not None:
CUDA_MAJOR = get_cuda_version() CUDA_MAJOR = get_cuda_version()
gencodes = ['-gencode', 'arch=compute_52,code=sm_52', gencodes = ['-gencode', 'arch=compute_50,code=sm_50',
'-gencode', 'arch=compute_52,code=sm_52',
'-gencode', 'arch=compute_60,code=sm_60', '-gencode', 'arch=compute_60,code=sm_60',
'-gencode', 'arch=compute_61,code=sm_61',] '-gencode', 'arch=compute_61,code=sm_61',]
......
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