PVD/metrics/ChamferDistancePytorch/chamfer2D/setup.py
Linqi (Alex) Zhou 2f6aa752a6 PVD
2021-10-19 13:54:46 -07:00

16 lines
459 B
Python

from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
setup(
name='chamfer_2D',
ext_modules=[
CUDAExtension('chamfer_2D', [
"/".join(__file__.split('/')[:-1] + ['chamfer_cuda.cpp']),
"/".join(__file__.split('/')[:-1] + ['chamfer2D.cu']),
]),
],
extra_cuda_cflags=['--compiler-bindir=/usr/bin/gcc-8'],
cmdclass={
'build_ext': BuildExtension
})