File size: 741 Bytes
29858c0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 |
# https://gist.github.com/tonyseek/7821993
import glob
import torch
from os import path as osp
from torch.utils.ffi import create_extension
abs_path = osp.dirname(osp.realpath(__file__))
extra_objects = [osp.join(abs_path, 'build/knn_cuda_kernel.so')]
extra_objects += glob.glob('/usr/local/cuda/lib64/*.a')
ffi = create_extension(
'knn_pytorch',
headers=['src/knn_pytorch.h'],
sources=['src/knn_pytorch.c'],
define_macros=[('WITH_CUDA', None)],
relative_to=__file__,
with_cuda=True,
extra_objects=extra_objects,
include_dirs=[osp.join(abs_path, 'include')]
)
if __name__ == '__main__':
assert torch.cuda.is_available(), 'Please install CUDA for GPU support.'
ffi.build()
|