Merge pull request #32 from 130ndim/master

Added flag for CPU install

setup.py

 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 @@ ext_modules = [
 ]
 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: