diff --git a/lib/make.sh b/lib/make.sh index 15a616bc..0b599f9a 100755 --- a/lib/make.sh +++ b/lib/make.sh @@ -1,4 +1,5 @@ TF_INC=$(python -c 'import tensorflow as tf; print(tf.sysconfig.get_include())') +NSYNC_INC=$TF_INC"/external/nsync/public" CUDA_PATH=/usr/local/cuda/ CXXFLAGS='' @@ -11,15 +12,15 @@ cd roi_pooling_layer if [ -d "$CUDA_PATH" ]; then nvcc -std=c++11 -c -o roi_pooling_op.cu.o roi_pooling_op_gpu.cu.cc \ - -I $TF_INC -D GOOGLE_CUDA=1 -x cu -Xcompiler -fPIC $CXXFLAGS \ + -I $TF_INC -I $NSYNC_INC -D GOOGLE_CUDA=1 -x cu -Xcompiler -fPIC $CXXFLAGS \ -arch=sm_37 g++ -std=c++11 -shared -o roi_pooling.so roi_pooling_op.cc \ - roi_pooling_op.cu.o -I $TF_INC -D GOOGLE_CUDA=1 -fPIC $CXXFLAGS \ + roi_pooling_op.cu.o -I $TF_INC -I $NSYNC_INC -D GOOGLE_CUDA=1 -fPIC $CXXFLAGS \ -lcudart -L $CUDA_PATH/lib64 else g++ -std=c++11 -shared -o roi_pooling.so roi_pooling_op.cc \ - -I $TF_INC -fPIC $CXXFLAGS + -I $TF_INC -I $NSYNC_INC -fPIC $CXXFLAGS fi cd ..