diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
index bfe744446a97b664adeae5569170b8128e6c52eb..ea27eba5130517c3d223d143f32a4c97a02ba491 100644
--- a/paddle/phi/infermeta/unary.cc
+++ b/paddle/phi/infermeta/unary.cc
@@ -4402,7 +4402,6 @@ void TraceInferMeta(
   auto sizes = vectorize(x_dims);
   if (x_dims.size() == 2) {
     sizes.clear();
-    sizes.push_back(1);
   } else {
     sizes.erase(sizes.begin() + std::max(dim1_, dim2_));
     sizes.erase(sizes.begin() + std::min(dim1_, dim2_));
diff --git a/paddle/phi/kernels/gpu/trace_kernel.cu b/paddle/phi/kernels/gpu/trace_kernel.cu
index 671ca490e136a2e5adb6e61eee7360f0ab9ff835..304bf778094d3de33b59b43ddaf7f886d2d4ceea 100644
--- a/paddle/phi/kernels/gpu/trace_kernel.cu
+++ b/paddle/phi/kernels/gpu/trace_kernel.cu
@@ -32,7 +32,10 @@ void TraceKernel(const Context& ctx,
   auto diag = funcs::Diagonal<T, Context>(ctx, &x, offset, axis1, axis2);
   if (diag.numel() > 0) {
     std::vector<int> reduce_dims;
-    reduce_dims.push_back(out->dims().size());
+    // Adapt to 0D output
+    auto out_dim_size = out->dims().size();
+    if (out_dim_size == 0) out_dim_size = 1;
+    reduce_dims.push_back(out_dim_size);
     funcs::ReduceKernel<T, T, kps::AddFunctor, kps::IdentityFunctor<T>>(
         ctx, diag, out, kps::IdentityFunctor<T>(), reduce_dims);
   } else {
diff --git a/paddle/phi/kernels/impl/trace_grad_kernel_impl.h b/paddle/phi/kernels/impl/trace_grad_kernel_impl.h
index 90a2327ef3e204fe0cda3cc281407926e0a61ba3..1099f27f3622e6156d9765f4af9dd60947de8714 100644
--- a/paddle/phi/kernels/impl/trace_grad_kernel_impl.h
+++ b/paddle/phi/kernels/impl/trace_grad_kernel_impl.h
@@ -91,7 +91,8 @@ void TraceGradKernel(const Context& ctx,
   auto input_dims = in_grad->dims();
   auto input_stride = phi::stride(input_dims);
   auto output_dims = out_grad.dims();
-  auto output_stride = phi::stride(output_dims);
+  auto output_stride = output_dims.size() == 0 ? phi::DDim(output_dims)
+                                               : phi::stride(output_dims);
 
   auto* out_data = out_grad.data<T>();
   T* x_data = ctx.template Alloc<T>(in_grad);
diff --git a/python/paddle/fluid/tests/unittests/test_zero_dim_tensor.py b/python/paddle/fluid/tests/unittests/test_zero_dim_tensor.py
index 7ea98f7c889a30a7b000ac329177b77c6a962730..965bcae57d9db6e4dfa0565c9198b525474e20d7 100644
--- a/python/paddle/fluid/tests/unittests/test_zero_dim_tensor.py
+++ b/python/paddle/fluid/tests/unittests/test_zero_dim_tensor.py
@@ -2437,6 +2437,16 @@ class TestSundryAPI(unittest.TestCase):
         self.assertEqual(b.grad.shape, [4, 5])
         self.assertEqual(c.grad.shape, [5])
 
+    def test_trace(self):
+        x = paddle.to_tensor([[3, 2], [1, 9]], dtype="float32")
+        x.stop_gradient = False
+        out = paddle.trace(x)
+        out.backward()
+
+        self.assertEqual(out.shape, [])
+        np.testing.assert_allclose(out, np.array(12))
+        self.assertEqual(x.grad.shape, [2, 2])
+
 
 class TestSundryAPIStatic(unittest.TestCase):
     def setUp(self):
@@ -4426,6 +4436,20 @@ class TestSundryAPIStatic(unittest.TestCase):
         self.assertEqual(res[2].shape, (4, 5))
         self.assertEqual(res[3].shape, (5,))
 
+    @prog_scope()
+    def test_trace(self):
+        x = paddle.to_tensor([[3, 2], [1, 9]], dtype="float32")
+        x.stop_gradient = False
+        out = paddle.trace(x)
+        paddle.static.append_backward(out)
+
+        prog = paddle.static.default_main_program()
+        res = self.exe.run(prog, fetch_list=[out, x.grad_name])
+
+        self.assertEqual(res[0].shape, ())
+        self.assertEqual(res[1].shape, (2, 2))
+        np.testing.assert_allclose(res[0], np.array(12))
+
 
 # Use to test API whose zero-dim input tensors don't have grad and not need to test backward in OpTest.
 class TestNoBackwardAPI(unittest.TestCase):