fixed unique install

de6a9fd7 · rusty1s · 8f8398e5 · de6a9fd7
Commit de6a9fd7 authored Nov 04, 2018 by rusty1s
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 1 deletion

setup.py setup.py +3 -1

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