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

Xreki · Jan 5, 2023 · ba1dce0 · ba1dce0
1 parent e8f4a32
commit ba1dce0
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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);