SoftGroup/lib/softgroup_ops/setup.py
2022-03-03 05:23:32 +00:00

15 lines
404 B
Python

from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
setup(
name='SOFTGROUP_OP',
ext_modules=[
CUDAExtension('SOFTGROUP_OP', [
'src/softgroup_api.cpp',
'src/softgroup_ops.cpp',
'src/cuda.cu'
], extra_compile_args={'cxx': ['-g'], 'nvcc': ['-O2']})
],
cmdclass={'build_ext': BuildExtension}
)