diff --git a/paddle/pten/kernels/gpu/trace_kernel.cu b/paddle/pten/kernels/gpu/trace_kernel.cu
index b0bfb3a8a51865df55a507506fe870016ac9181e..155bfbd02af17c390e83a84a269df01a77b11637 100644
--- a/paddle/pten/kernels/gpu/trace_kernel.cu
+++ b/paddle/pten/kernels/gpu/trace_kernel.cu
@@ -33,9 +33,8 @@ void TraceKernel(const Context& ctx,
     auto stream = ctx.stream();
     std::vector<int> reduce_dims;
     reduce_dims.push_back(out->dims().size());
-    kernels::
-        TensorReduceFunctorImpl<T, T, kps::AddFunctor, kps::IdentityFunctor<T>>(
-            ctx, diag, out, kps::IdentityFunctor<T>(), reduce_dims, stream);
+    kernels::TensorReduceImpl<T, T, kps::AddFunctor, kps::IdentityFunctor<T>>(
+        ctx, diag, out, kps::IdentityFunctor<T>(), reduce_dims, stream);
   } else {
     paddle::operators::math::SetConstant<Context, T> functor;
     functor(ctx, out, static_cast<T>(0));