build.sh 352 Bytes
Newer Older
rusty1s's avatar
rusty1s committed
1
2
#!/bin/sh

rusty1s's avatar
rusty1s committed
3
4
echo "Compiling kernel..."

rusty1s's avatar
rusty1s committed
5
6
7
8
9
TORCH=$(python -c "import os; import torch; print(os.path.dirname(torch.__file__))")
SRC_DIR=torch_scatter/kernel
BUILD_DIR=torch_scatter/build

mkdir -p $BUILD_DIR
rusty1s's avatar
rusty1s committed
10
nvcc -c -o $BUILD_DIR/kernel.so $SRC_DIR/kernel.cu -arch=sm_30 -Xcompiler -fPIC -shared -I$TORCH/lib/include/TH -I$TORCH/lib/include/THC -I$SRC_DIR