Fix

Author: co63oc

paddle/phi/kernels/cpu/p_norm_grad_kernel.cc

@@ -59,6 +59,7 @@ void PNormGradKernel(const Context& dev_ctx,
   auto* in_norm_dy = &out_grad;
   auto* out_dx = x_grad;
   dev_ctx.template Alloc<T>(out_dx);
+  if (x.numel() == 0) return;
 
   T eps = static_cast<T>(epsilon);
   auto xdim = in_x->dims();

paddle/phi/kernels/cpu/p_norm_kernel.cc

@@ -63,10 +63,11 @@ void PNormKernel(const Context& dev_ctx,
   GetDims(xdim, axis, &pre, &n, &post, asvector);
 
   if (x.numel() == 0) {
-    if (out->numel() > 0) {
-      std::vector<int64_t> vec_dims = common::vectorize(out->dims());
-      phi::Full<T, Context>(
-          dev_ctx, phi::IntArray(vec_dims), static_cast<T>(0), out);
+    if (out) {
+      phi::Full<T, Context>(dev_ctx,
+                            phi::IntArray(common::vectorize(out->dims())),
+                            static_cast<T>(0),
+                            out);
     }
     return;
   }

paddle/phi/kernels/gpu/p_norm_grad_kernel.cu

@@ -83,6 +83,7 @@ void PNormGradKernel(const Context& dev_ctx,
   auto* in_norm_dy = &out_grad;
   auto* out_dx = x_grad;
   dev_ctx.template Alloc<T>(out_dx);
+  if (x.numel() == 0) return;
 
   auto xdim = in_x->dims();
   bool reduce_all = (in_norm->numel() == 1);

paddle/phi/kernels/gpu/p_norm_kernel.cu

@@ -95,10 +95,11 @@ void PNormKernel(const Context& dev_ctx,
       funcs::details::GetReduceDim(axis_dims, xdim.size(), asvector);
 
   if (x.numel() == 0) {
-    if (out->numel() > 0) {
-      std::vector<int64_t> vec_dims = common::vectorize(out->dims());
-      phi::Full<T, Context>(
-          dev_ctx, phi::IntArray(vec_dims), static_cast<T>(0), out);
+    if (out) {
+      phi::Full<T, Context>(dev_ctx,
+                            phi::IntArray(common::vectorize(out->dims())),
+                            static_cast<T>(0),
+                            out);
     }
     return;
   }

test/legacy_test/test_norm_op.py

@@ -212,6 +212,13 @@ def test_norm_x_type():
             self.assertRaises(TypeError, test_norm_x_type)
 
 
+class TestNormOp_ZeroSize(TestNormOp):
+    def init_test_case(self):
+        self.shape = [5, 3, 0, 7]
+        self.axis = 0
+        self.epsilon = 1e-8
+
+
 if __name__ == '__main__':
     paddle.enable_static()
     unittest.main()