From 6838dd4f0ff51b3e60168752984b3a7cb21e7269 Mon Sep 17 00:00:00 2001 From: Anjali Sridhar Date: Sun, 11 Apr 2021 16:29:38 +0000 Subject: [PATCH] use exported path --- setup.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/setup.py b/setup.py index 8e16a06d69..198240b0fa 100644 --- a/setup.py +++ b/setup.py @@ -794,6 +794,13 @@ def is_torch_cuda(build_ext, include_dirs, extra_compile_args): def build_torch_extension(build_ext, options, torch_version): pytorch_compile_flags = ["-std=c++14" if flag == "-std=c++11" else flag for flag in options['COMPILE_FLAGS']] + cuda_include_dirs, cuda_lib_dirs = get_cuda_dirs( + build_ext, options['COMPILE_FLAGS']) + options['INCLUDES'] += cuda_include_dirs + options['LIBRARY_DIRS'] += cuda_lib_dirs + # options['MACROS'] += [('HAVE_CUDA', '1')] + # options['LIBRARIES'] += ['cudart'] + have_cuda = is_torch_cuda(build_ext, include_dirs=options['INCLUDES'], extra_compile_args=pytorch_compile_flags) if not have_cuda and check_macro(options['MACROS'], 'HAVE_CUDA'):