Commit 33e81aa6 authored by Shagun Sodhani's avatar Shagun Sodhani
Browse files

updated setup.py as per @rusty1s's suggestion

parent 92082f98
...@@ -13,13 +13,10 @@ cmdclass = {} ...@@ -13,13 +13,10 @@ cmdclass = {}
if torch.cuda.is_available(): if torch.cuda.is_available():
ext_modules += [ 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=['-l', 'cusparse']
extra_link_args=['-lcusparse'], )
),
CUDAExtension('unique_cuda',
['cuda/unique.cpp', 'cuda/unique_kernel.cu'])
] ]
cmdclass['build_ext'] = BuildExtension cmdclass['build_ext'] = BuildExtension
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment