File size: 392 Bytes
33569f9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 | import torch
from setuptools import setup, Extension
from torch.utils.cpp_extension import BuildExtension, CppExtension
setup(
name='nms_1d_cpu',
ext_modules=[
CppExtension(
name = 'nms_1d_cpu',
sources = ['./csrc/nms_cpu.cpp'],
extra_compile_args=['-fopenmp']
)
],
cmdclass={
'build_ext': BuildExtension
}
)
|