fix div 0 bug (#55644)

paddle/phi/kernels/cpu/nanmedian_kernel.cc

@@ -116,6 +116,14 @@ void ProcessMedianKernel(const Context& dev_ctx,
   auto x_dim = x.dims();
   int64_t x_rank = x_dim.size();
   int64_t stride = x_dim[x_rank - 1];
+
+  PADDLE_ENFORCE_NE(
+      stride,
+      0,
+      phi::errors::InvalidArgument("The input Tensor x's shape[-1] should not "
+                                   "be 0, but shape is %s now.",
+                                   x_dim));
+
   int64_t pre_dim = numel / stride;
   int64_t i = 0;
 

paddle/phi/kernels/gpu/matrix_rank_tol_kernel.cu

@@ -340,6 +340,18 @@ void MatrixRankTolKernel(const Context& dev_ctx,
   auto dim_out = out->dims();
   int rows = dim_x[dim_x.size() - 2];
   int cols = dim_x[dim_x.size() - 1];
+  PADDLE_ENFORCE_NE(
+      rows,
+      0,
+      phi::errors::InvalidArgument("The input Tensor x's shape[-2] should not "
+                                   "be 0, but shape is %s now.",
+                                   dim_x));
+  PADDLE_ENFORCE_NE(
+      cols,
+      0,
+      phi::errors::InvalidArgument("The input Tensor x's shape[-1] should not "
+                                   "be 0, but shape is %s now.",
+                                   dim_x));
   int k = std::min(rows, cols);
   auto numel = x.numel();
   int batches = numel / (rows * cols);

paddle/phi/kernels/gpu/nanmedian_kernel.cu

@@ -149,6 +149,14 @@ void ProcessMedianKernel(const Context& dev_ctx,
   auto x_dim = x.dims();
   int64_t x_rank = x_dim.size();
   int64_t stride = x_dim[x_rank - 1];
+
+  PADDLE_ENFORCE_NE(
+      stride,
+      0,
+      phi::errors::InvalidArgument("The input Tensor x's shape[-1] should not "
+                                   "be 0, but shape is %s now.",
+                                   x_dim));
+
   int64_t pre_dim = numel / stride;
   int64_t i = 0;