Merge pull request #2 from shagunsodhani/master

updated setup.py as per @rusty1s's suggestion

Merge pull request #2 from shagunsodhani/master
updated setup.py as per @rusty1s's suggestion
85c6c492 · Matthias Fey · GitHub · 43c50459 · 00566354 · 85c6c492
Unverified Commit 85c6c492 authored Oct 20, 2018 by Matthias Fey Committed by GitHub Oct 20, 2018
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 7 deletions

setup.py setup.py +4 -7

No files found.
--- a/setup.py
+++ b/setup.py
@@ -13,13 +13,10 @@ cmdclass = {}
 if torch.cuda.is_available():
    ext_modules += [
-        CUDAExtension(
+        CUDAExtension('spspmm_cuda',
-            'spspmm_cuda',
+                      ['cuda/spspmm.cpp', 'cuda/spspmm_kernel.cu'],
-            ['cuda/spspmm.cpp', 'cuda/spspmm_kernel.cu'],
+                      extra_link_args=['-lcusparse', '-l', 'cusparse'],
-            extra_link_args=['-lcusparse'],
+                      )
-        ),
-        CUDAExtension('unique_cuda',
-                      ['cuda/unique.cpp', 'cuda/unique_kernel.cu'])
    ]
    cmdclass['build_ext'] = BuildExtension