From dce05da5350962fd2976eeca830b84a8bf29e29d Mon Sep 17 00:00:00 2001
From: Jun Gao <imgaojun@gmail.com>
Date: Tue, 21 Nov 2023 15:09:35 +0800
Subject: [PATCH] fix thrust-transform-reduce error (#5078)

---
 colossalai/kernel/cuda_native/csrc/kernels/cuda_util.cu | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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 */