Merge pull request #32 from 130ndim/master

Added flag for CPU install

Author: rusty1s

setup.py

@@ -1,5 +1,6 @@
 import platform
 from setuptools import setup, find_packages
+from sys import argv
 import torch
 from torch.utils.cpp_extension import CppExtension, CUDAExtension, CUDA_HOME
 
@@ -16,7 +17,13 @@
 ]
 cmdclass = {'build_ext': torch.utils.cpp_extension.BuildExtension}
 
-if CUDA_HOME is not None:
+GPU = True
+for arg in argv:
+    if arg == '--cpu':
+        GPU = False
+        argv.remove(arg)
+
+if CUDA_HOME is not None and GPU:
     if platform.system() == 'Windows':
         extra_link_args = ['cusparse.lib']
     else: