Fix

Author: co63oc

paddle/phi/kernels/gpudnn/softmax_grad_kernel.cu

@@ -28,6 +28,7 @@ void SoftmaxGradGPUDNNKernel(const Context& dev_ctx,
                              int axis,
                              DenseTensor* x_grad) {
   dev_ctx.template Alloc<T>(x_grad);
+  if (x_grad->numel() == 0) return;
 
   const int rank = out.dims().size();
   // For 0D Tensor

paddle/phi/kernels/gpudnn/softmax_kernel.cu

@@ -27,6 +27,7 @@ void SoftmaxGPUDNNKernel(const Context& dev_ctx,
                          int axis,
                          DenseTensor* out) {
   dev_ctx.template Alloc<T>(out);
+  if (x.numel() == 0) return;
 
   const int rank = x.dims().size();
   // For 0D Tensor

test/legacy_test/test_softmax_op.py

@@ -15,7 +15,7 @@
 import unittest
 
 import numpy as np
-from op_test import OpTest, convert_float_to_uint16
+from op_test import OpTest, convert_float_to_uint16, get_places
 from utils import static_guard
 
 import paddle
@@ -646,5 +646,20 @@ def executed_api(self):
         self.softmax = F.softmax_
 
 
+class TestSoftmaxAPI_ZeroSize(unittest.TestCase):
+    def test_dygraph(self):
+        for place in get_places():
+            paddle.disable_static(place)
+            x = paddle.rand([0, 2, 3])
+            x.stop_gradient = False
+            x.retain_grads()
+            out = paddle.nn.functional.softmax(x)
+            out.retain_grads()
+            out.backward()
+            np.testing.assert_allclose(out.numpy(), np.random.random([0, 2, 3]))
+            np.testing.assert_allclose(x.grad.shape, x.shape)
+            paddle.enable_static()
+
+
 if __name__ == "__main__":
     unittest.main()