Skip to content

GitLab

Commit 5b2cf9ca authored by rusty1s's avatar rusty1s
Browse files

only include kernel if cuda available

parent 2db5d1f3
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
build.py
View file @ 5b2cf9ca
...@@ -5,7 +5,7 @@ from torch.utils.ffi import create_extension ...@@ -5,7 +5,7 @@ from torch.utils.ffi import create_extension
headers = ['torch_scatter/src/cpu.h'] headers = ['torch_scatter/src/cpu.h']
sources = ['torch_scatter/src/cpu.c'] sources = ['torch_scatter/src/cpu.c']
include_dirs = ['torch_scatter/src', 'torch_scatter/kernel'] include_dirs = ['torch_scatter/src']
define_macros = [] define_macros = []
extra_objects = [] extra_objects = []
with_cuda = False with_cuda = False
...@@ -15,6 +15,7 @@ if torch.cuda.is_available(): ...@@ -15,6 +15,7 @@ if torch.cuda.is_available():
headers += ['torch_scatter/src/cuda.h'] headers += ['torch_scatter/src/cuda.h']
sources += ['torch_scatter/src/cuda.c'] sources += ['torch_scatter/src/cuda.c']
include_dirs += ['torch_scatter/kernel']
define_macros += [('WITH_CUDA', None)] define_macros += [('WITH_CUDA', None)]
extra_objects += ['torch_scatter/build/kernel.so'] extra_objects += ['torch_scatter/build/kernel.so']
with_cuda = True with_cuda = True
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment