fix log_softmax if any dimension is 0-d (#34635)

PaddlePaddle · Aug 6, 2021 · 436a9f1 · 436a9f1
1 parent 06651c4
commit 436a9f1
Showing 1 changed file with 9 additions and 4 deletions.
diff --git a/paddle/fluid/operators/log_softmax_op.h b/paddle/fluid/operators/log_softmax_op.h
@@ -131,8 +131,10 @@ class LogSoftmaxKernel : public framework::OpKernel<T> {
     // allocate memory on device.
     Out->mutable_data<T>(context.GetPlace());
 
-    LogSoftmaxFunctor<DeviceContext, T>()(
-        context.template device_context<DeviceContext>(), X, Out, axis);
+    if (X->numel() != 0) {
+      LogSoftmaxFunctor<DeviceContext, T>()(
+          context.template device_context<DeviceContext>(), X, Out, axis);
+    }
   }
 };
 
@@ -183,8 +185,11 @@ class LogSoftmaxGradKernel : public framework::OpKernel<T> {
     // allocate memory on device.
     dX->mutable_data<T>(context.GetPlace());
 
-    LogSoftmaxGradFunctor<DeviceContext, T>()(
-        context.template device_context<DeviceContext>(), Out, dOut, dX, axis);
+    if (Out->numel() != 0) {
+      LogSoftmaxGradFunctor<DeviceContext, T>()(
+          context.template device_context<DeviceContext>(), Out, dOut, dX,
+          axis);
+    }
   }
 };