pass torch path

build.py

-import sys
-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_spline_conv/src/cpu.h']
@@ -17,7 +16,7 @@ extra_objects = []
 with_cuda = False
 
 if torch.cuda.is_available():
-    subprocess.call('./build.sh {}'.format(sys.executable))
+    subprocess.call('./build.sh {}'.format(osp.dirname(torch.__file__)))
 
     headers += ['torch_spline_conv/src/cuda.h']
     sources += ['torch_spline_conv/src/cuda.c']

build.sh

@@ -2,10 +2,8 @@
 
 echo "Compiling kernel..."
 
-PYTHON=${1:-python}
-TORCH=$($PYTHON -c "import os; import torch; print(os.path.dirname(torch.__file__))")
 SRC_DIR=torch_spline_conv/kernel
 BUILD_DIR=torch_spline_conv/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
+mkdir -p "$BUILD_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"