From ba1dce0a1af618add4063d9c91ee22e2fc4a8c0c Mon Sep 17 00:00:00 2001
From: Zhong Hui <zhonghui.net@gmail.com>
Date: Thu, 5 Jan 2023 14:17:04 +0800
Subject: [PATCH] [BugFix] Fix illegal memory overflow for p_norm op (#49537)

---
 paddle/phi/infermeta/unary.cc | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
index a4171825a36..998001a008e 100644
--- 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);
-- 
GitLab