adding cusparse as compile arg

afccb12e · rusty1s · 5836042b · afccb12e
Commit afccb12e authored Aug 20, 2018 by rusty1s
Hide whitespace changes
Inline Side-by-side

Showing with 5 additions and 2 deletions

setup.py setup.py +5 -2

No files found.
--- a/setup.py
+++ b/setup.py
@@ -13,8 +13,11 @@ cmdclass = {}
 if torch.cuda.is_available():
    ext_modules += [
-        CUDAExtension('spspmm_cuda',
+        CUDAExtension(
-                      ['cuda/spspmm.cpp', 'cuda/spspmm_kernel.cu'])
+            'spspmm_cuda', ['cuda/spspmm.cpp', 'cuda/spspmm_kernel.cu'],
+            extra_compile_args={
+                'nvcc': ['-lcusparse']
+            })
    ]
    cmdclass['build_ext'] = BuildExtension