rename

build.py

@@ -18,7 +18,7 @@ extra_objects = []
 with_cuda = False
 
 if torch.cuda.is_available():
-    subprocess.call(['./build_new.sh', osp.dirname(torch.__file__)])
+    subprocess.call(['./build.sh', osp.dirname(torch.__file__)])
 
     headers += ['aten/THCC/THCC{}.h'.format(f) for f in files]
     sources += ['aten/THCC/THCC{}.c'.format(f) for f in files]

build.sh

@@ -3,10 +3,8 @@
 echo "Compiling kernel..."
 
 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
+SRC_DIR=aten/THC
+BUILD_DIR=aten/build
 
 mkdir -p $BUILD_DIR
-for i in serial grid; do
-  $(which nvcc) -c -o "$BUILD_DIR/$i.so" "$SRC_DIR/$i.cu" -arch=sm_35 -Xcompiler -fPIC -shared "-I$TORCH/lib/include/TH" "-I$TORCH/lib/include/THC" "-I$SRC_DIR"
-done
+$(which nvcc) -c -o "$BUILD_DIR/THC.so" "$SRC_DIR/THC.cu" -arch=sm_52 -Xcompiler -fPIC -shared "-I$TORCH/lib/include/TH" "-I$TORCH/lib/include" "-I$SRC_DIR"

build_new.sh

-#!/bin/sh
-
-echo "Compiling kernel..."
-
-if [ -z "$1" ]; then TORCH=$(python -c "import os; import torch; print(os.path.dirname(torch.__file__))"); else TORCH="$1"; fi
-SRC_DIR=aten/THC
-BUILD_DIR=aten/build
-
-mkdir -p $BUILD_DIR
-$(which nvcc) -c -o "$BUILD_DIR/THC.so" "$SRC_DIR/THC.cu" -arch=sm_52 -Xcompiler -fPIC -shared "-I$TORCH/lib/include/TH" "-I$TORCH/lib/include" "-I$SRC_DIR"