Skip to content

GitLab

Commit b51c22a4 authored by rusty1s's avatar rusty1s
Browse files

USE_GPU

parent 80a7dc52
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 5 deletions
setup.py
View file @ b51c22a4
...@@ -7,10 +7,6 @@ from sys import argv ...@@ -7,10 +7,6 @@ from sys import argv
import torch import torch
from torch.utils.cpp_extension import CppExtension, CUDAExtension, CUDA_HOME from torch.utils.cpp_extension import CppExtension, CUDAExtension, CUDA_HOME
USE_GPU = True
if '--cpu' in argv:
USE_GPU = False
cxx_extra_compile_args = [] cxx_extra_compile_args = []
nvcc_extra_compile_args = ['-arch=sm_35', '--expt-relaxed-constexpr'] nvcc_extra_compile_args = ['-arch=sm_35', '--expt-relaxed-constexpr']
if platform.system() != 'Windows': if platform.system() != 'Windows':
...@@ -30,7 +26,7 @@ ext_modules += [ ...@@ -30,7 +26,7 @@ ext_modules += [
extra_compile_args=cxx_extra_compile_args) for ext in exts extra_compile_args=cxx_extra_compile_args) for ext in exts
] ]
if CUDA_HOME is not None and USE_GPU: if CUDA_HOME is not None and '--cpu' not in argv:
exts = [e.split(osp.sep)[-1][:-4] for e in glob(osp.join('cuda', '*.cpp'))] exts = [e.split(osp.sep)[-1][:-4] for e in glob(osp.join('cuda', '*.cpp'))]
ext_modules += [ ext_modules += [
CUDAExtension( CUDAExtension(
......
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