From 389023ff318ca7b4925518bc34af62a286b676ff Mon Sep 17 00:00:00 2001 From: ver217 Date: Wed, 13 Nov 2024 10:32:30 +0800 Subject: [PATCH] [hotfix] fix include path --- setup.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/setup.py b/setup.py index 3c23286..19bd9ca 100644 --- a/setup.py +++ b/setup.py @@ -20,6 +20,7 @@ def check_uring_compatibility(): kernel_version = version.parse(uname_info.release.split("-")[0]) return kernel_version >= version.parse("5.10") + def check_pthread_compatibility(): uname_info = uname() if uname_info.system != "Linux": @@ -39,7 +40,7 @@ def check_pthread_compatibility(): if os.environ.get("DISABLE_AIO") == "1": enable_aio = False if os.environ.get("DISABLE_PTHREAD") == "1" or not check_pthread_compatibility(): - enable_pthread=False + enable_pthread = False assert enable_aio or enable_uring or enable_pthread if os.environ.get("WITH_ROOT") == "1": @@ -65,7 +66,12 @@ def check_pthread_compatibility(): def cpp_ext_helper(name, sources, **kwargs): - from torch.utils.cpp_extension import CppExtension + from torch.utils.cpp_extension import CUDA_HOME, CppExtension + + extra_include_dirs = [] + + if CUDA_HOME is not None: + extra_include_dirs.append(os.path.join(CUDA_HOME, "include")) extra_include_dirs = [] if "C_INCLUDE_PATH" in os.environ: