Merge pull request #12 from 130ndim/master

Added flag for CPU install

setup.py

 from setuptools import setup, find_packages
+from sys import argv
 import torch
 from torch.utils.cpp_extension import CppExtension, CUDAExtension, CUDA_HOME
 
@@ -17,7 +18,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:
     ext_modules += [
         CUDAExtension('torch_spline_conv.basis_cuda',
                       ['cuda/basis.cpp', 'cuda/basis_kernel.cu'],