fix(torch): add include sometimes needed

jolibrain · Sep 15, 2021 · 74487dc · 74487dc
1 parent 67e5bc0
commit 74487dc
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -747,6 +747,7 @@ if (USE_TORCH)
     ${PYTORCH_PATCHES_PATH}/pytorch_15_compile.patch
     ${PYTORCH_PATCHES_PATH}/pytorch_19_new_logger.patch
     ${PYTORCH_PATCHES_PATH}/pytorch_19_use_new_logger.patch
+    ${PYTORCH_PATCHES_PATH}/pytorch_19_thrust_include.patch
   )
 
   message(STATUS "Configuring libtorch")

diff --git a/patches/pytorch/pytorch_19_thrust_include.patch b/patches/pytorch/pytorch_19_thrust_include.patch
@@ -0,0 +1,12 @@
+diff --git a/caffe2/utils/math_gpu.cu b/caffe2/utils/math_gpu.cu
+index 3de3ace02e..0b407ab545 100644
+--- a/caffe2/utils/math_gpu.cu
++++ b/caffe2/utils/math_gpu.cu
+@@ -10,6 +10,7 @@
+ #include <cub/block/block_reduce.cuh>
+ #include <cub/cub.cuh>
+
++#include <thrust/host_vector.h>
+ #include <thrust/device_vector.h>
+ #include <thrust/functional.h>
+