diff --git a/colossalai/kernel/cuda_native/csrc/kernels/cuda_util.cu b/colossalai/kernel/cuda_native/csrc/kernels/cuda_util.cu
index 9a6a8ebc3..e5ac17308 100644
--- a/colossalai/kernel/cuda_native/csrc/kernels/cuda_util.cu
+++ b/colossalai/kernel/cuda_native/csrc/kernels/cuda_util.cu
@@ -1,6 +1,6 @@
 #include <thrust/device_vector.h>
 #include <thrust/reduce.h>
-
+#include <thrust/transform_reduce.h>
 #include "cuda_util.h"
 
 /* GPU function guard */