[BugFix] Fix illegal memory overflow for p_norm op (#49537)

ba1dce0a · Zhong Hui · GitHub · e8f4a327 · ba1dce0a
隐藏空白更改
内联并排

Showing with 2 addition and 1 deletion

paddle/phi/infermeta/unary.cc paddle/phi/infermeta/unary.cc +2 -1

未找到文件。
--- a/paddle/phi/infermeta/unary.cc
+++ b/paddle/phi/infermeta/unary.cc
@@ -2644,8 +2644,9 @@ void PNormInferMeta(const MetaTensor& x,
    if (reduce_dims.size() == 0) {
      reduce_dims.emplace_back(1);
    }
+
+    x_dim[axis] = 1;
  }
-  x_dim[axis] = 1;

  if (keepdim) {
    out->set_dims(x_dim);