Fix build script compatibility with Pytorch v1.5+

mapillary · Nov 5, 2020 · 3d10aea · 3d10aea
1 parent 574e89a
commit 3d10aea
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/setup.py b/setup.py
@@ -3,6 +3,8 @@
 import setuptools
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension
 
+here = path.abspath(path.dirname(__file__))
+
 
 def find_sources(root_dir):
     sources = []
@@ -22,12 +24,10 @@ def make_extension(name, package):
             "cxx": ["-O3"],
             "nvcc": ["--expt-extended-lambda"],
         },
-        include_dirs=["include/"],
+        include_dirs=[path.join(here, "include")],
     )
 
 
-here = path.abspath(path.dirname(__file__))
-
 with open(path.join(here, "README.md"), encoding="utf-8") as f:
     long_description = f.read()