fix trace offset out of shape (#33922)

paddle/fluid/operators/trace_op.cu

@@ -14,6 +14,7 @@
 
 #include <thrust/device_vector.h>
 #include <thrust/host_vector.h>
+#include "paddle/fluid/operators/math/math_function.h"
 #include "paddle/fluid/operators/reduce_ops/cub_reduce.h"
 #include "paddle/fluid/operators/trace_op.h"
 
@@ -50,6 +51,9 @@ class TraceCUDAKernel : public framework::OpKernel<T> {
       TensorReduce<T, T, cub::Sum, IdentityFunctor>(
           diag, out, reduce_dims, static_cast<T>(0), cub::Sum(),
           IdentityFunctor(), stream);
+    } else {
+      math::SetConstant<DeviceContext, T> functor;
+      functor(context.device_context<DeviceContext>(), out, static_cast<T>(0));
     }
   }
 };

paddle/fluid/operators/trace_op.h

@@ -179,7 +179,7 @@ class TraceKernel : public framework::OpKernel<T> {
 
     auto output_dims = out->dims();
 
-    out->mutable_data<T>(context.GetPlace());
+    T* out_data = out->mutable_data<T>(context.GetPlace());
 
     const framework::Tensor diag =
         Diagonal<DeviceContext, T>(context, input, offset, dim1, dim2);
@@ -191,6 +191,8 @@ class TraceKernel : public framework::OpKernel<T> {
       auto reduce_dim = Eigen::array<int, 1>({1});
       output.device(place) = x.sum(reduce_dim);
       out->Resize(output_dims);
+    } else {
+      std::fill(out_data, out_data + out->numel(), static_cast<T>(0));
     }
   }
 };