23 lines
711 B
Python
23 lines
711 B
Python
|
from setuptools import setup
|
||
|
|
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
|
||
|
|
|
||
|
|
|
||
|
|
def make_cuda_ext(name, sources):
|
||
|
|
|
||
|
|
return CUDAExtension(
|
||
|
|
name='{}'.format(name), sources=[p for p in sources], extra_compile_args={
|
||
|
|
'cxx': [],
|
||
|
|
'nvcc': [
|
||
|
|
'-D__CUDA_NO_HALF_OPERATORS__',
|
||
|
|
'-D__CUDA_NO_HALF_CONVERSIONS__',
|
||
|
|
'-D__CUDA_NO_HALF2_OPERATORS__',
|
||
|
|
]
|
||
|
|
})
|
||
|
|
|
||
|
|
|
||
|
|
setup(
|
||
|
|
name='deform_conv', ext_modules=[
|
||
|
|
make_cuda_ext(name='deform_conv_cuda',
|
||
|
|
sources=['src/deform_conv_cuda.cpp', 'src/deform_conv_cuda_kernel.cu'])
|
||
|
|
], cmdclass={'build_ext': BuildExtension}, zip_safe=False)
|