SoftGroup/softgroup/lib/softgroup_ops/setup.py
2022-04-08 11:58:58 +00:00

16 lines
429 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})