fixed kernel build

build.py

-import os
+import os.path as osp
 import shutil
 import subprocess
 
 import torch
 from torch.utils.ffi import create_extension
 
-if os.path.exists('build'):
+if osp.exists('build'):
     shutil.rmtree('build')
 
 headers = ['torch_cluster/src/cpu.h']
@@ -16,7 +16,7 @@ extra_objects = []
 with_cuda = False
 
 if torch.cuda.is_available():
-    subprocess.call('./build.sh')
+    subprocess.call(['./build.sh', osp.dirname(torch.__file__)])
 
     headers += ['torch_cluster/src/cuda.h']
     sources += ['torch_cluster/src/cuda.c']

build.sh

@@ -2,9 +2,9 @@
 
 echo "Compiling kernel..."
 
-TORCH=$(python -c "import os; import torch; print(os.path.dirname(torch.__file__))")
+if [ -z "$1" ]; then TORCH=$(python -c "import os; import torch; print(os.path.dirname(torch.__file__))"); else TORCH="$1"; fi
 SRC_DIR=torch_cluster/kernel
 BUILD_DIR=torch_cluster/build
 
 mkdir -p $BUILD_DIR
-$(which nvcc) -c -o $BUILD_DIR/kernel.so $SRC_DIR/kernel.cu -arch=sm_35 -Xcompiler -fPIC -shared -I$TORCH/lib/include/TH -I$TORCH/lib/include/THC -I$SRC_DIR
+$(which nvcc) -c -o "$BUILD_DIR/kernel.so" "$SRC_DIR/kernel.cu" -arch=sm_35 -Xcompiler -fPIC -shared "-I$1/lib/include/TH" "-I$1/lib/include/THC" "-I$SRC_DIR"