diff --git a/setup.py b/setup.py index ed4720c3..1cb2ec59 100644 --- a/setup.py +++ b/setup.py @@ -4,6 +4,7 @@ import shutil import platform import subprocess +import sysconfig from setuptools import setup, Extension from setuptools.command.build_ext import build_ext @@ -38,10 +39,12 @@ def run(self): def build_extension(self, ext): extdir = os.path.abspath(os.path.dirname(self.get_ext_fullpath(ext.name))) + python_include_dir = sysconfig.get_path("include") cmake_args = [ "-DCMAKE_CUDA_COMPILER=nvcc", "-DCMAKE_LIBRARY_OUTPUT_DIRECTORY=" + extdir, "-DPYTHON_EXECUTABLE=" + sys.executable, + "-DPYTHON_INCLUDE_DIR=" + python_include_dir, ] cfg = "Debug" if self.debug else "Release"