setup.py

from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
import glob
import sys

# build clib
src_root = "clib"
sources = glob.glob(f"{src_root}/src/*.cpp") + glob.glob(f"{src_root}/src/*.cu")
includes = f"{sys.path[0]}/{src_root}/include"

setup(
    name='dvs',
    ext_modules=[
        CUDAExtension(
            name='clib._ext',
            sources=sources,
            extra_compile_args={
                "cxx": ["-O2", f"-I{includes}"],
                "nvcc": ["-O2", f"-I{includes}"],
            },
        )
    ],
    cmdclass={
        'build_ext': BuildExtension
    }
)