cppextension and fixed include_dirs

meder411 · meder411 · commit 9b6473fe95c5 · 2020-05-06T19:11:19.000-04:00
diff --git a/demo/setup_ninja.py b/demo/setup_ninja.py
@@ -1,15 +1,17 @@
 from setuptools import setup
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+from torch.utils.cpp_extension import BuildExtension, CppExtension
 
 setup(
     name='lltm_cuda',
     ext_modules=[
-        CUDAExtension('lltm_cuda', [
-            'lltm_cuda.cpp',
-            'lltm_cuda_kernel.cu',
-        ]),
+        CppExtension(
+            'lltm_cuda',
+            [
+                'lltm_cuda.cpp',
+            ],
+            include_dirs=['/abs/path/to/includes'],
+        ),
     ],
-    include_dirs=['includes'],
     cmdclass={
         'build_ext': BuildExtension
     })
diff --git a/demo/setup_no_ninja.py b/demo/setup_no_ninja.py
@@ -1,15 +1,17 @@
 from setuptools import setup
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+from torch.utils.cpp_extension import BuildExtension, CppExtension
 
 setup(
     name='lltm_cuda',
     ext_modules=[
-        CUDAExtension('lltm_cuda', [
-            'lltm_cuda.cpp',
-            'lltm_cuda_kernel.cu',
-        ]),
+        CppExtension(
+            'lltm_cuda',
+            [
+                'lltm_cuda.cpp',
+            ],
+            include_dirs=['includes'],
+        ),
     ],
-    include_dirs=['includes'],
     cmdclass={
         'build_ext': BuildExtension.with_options(use_ninja=False)
     })
