From 628ff34b6c6070bd40f0240566c900483f160f4d Mon Sep 17 00:00:00 2001
From: whs <wanghaoshuang@baidu.com>
Date: Sun, 26 Sep 2021 13:41:03 +0800
Subject: [PATCH] Fix FPE of label smooth op (#35861)

---
 paddle/fluid/operators/label_smooth_op.h | 43 +++++++++++++-----------
 1 file changed, 23 insertions(+), 20 deletions(-)

diff --git a/paddle/fluid/operators/label_smooth_op.h b/paddle/fluid/operators/label_smooth_op.h
index 760d542505e..6b509eb64cc 100644
--- a/paddle/fluid/operators/label_smooth_op.h
+++ b/paddle/fluid/operators/label_smooth_op.h
@@ -29,20 +29,21 @@ class LabelSmoothKernel : public framework::OpKernel<T> {
     auto* dist_t = ctx.Input<framework::Tensor>("PriorDist");
     auto label_dim = in_t->dims()[in_t->dims().size() - 1];
     out_t->mutable_data<T>(ctx.GetPlace());
-
-    auto epsilon = ctx.Attr<float>("epsilon");
-    auto out = framework::EigenVector<T>::Flatten(*out_t);
-    auto in = framework::EigenVector<T>::Flatten(*in_t);
-    auto& dev = *ctx.template device_context<DeviceContext>().eigen_device();
-    if (dist_t) {
-      auto dist = framework::EigenVector<T>::Flatten(*dist_t);
-      out.device(dev) =
-          static_cast<T>(1 - epsilon) * in +
-          static_cast<T>(epsilon) *
-              dist.broadcast(Eigen::DSizes<int, 1>(in_t->numel() / label_dim));
-    } else {
-      out.device(dev) = static_cast<T>(1 - epsilon) * in +
-                        static_cast<T>(epsilon / label_dim);
+    if (label_dim != 0) {
+      auto epsilon = ctx.Attr<float>("epsilon");
+      auto out = framework::EigenVector<T>::Flatten(*out_t);
+      auto in = framework::EigenVector<T>::Flatten(*in_t);
+      auto& dev = *ctx.template device_context<DeviceContext>().eigen_device();
+      if (dist_t) {
+        auto dist = framework::EigenVector<T>::Flatten(*dist_t);
+        out.device(dev) = static_cast<T>(1 - epsilon) * in +
+                          static_cast<T>(epsilon) *
+                              dist.broadcast(Eigen::DSizes<int, 1>(
+                                  in_t->numel() / label_dim));
+      } else {
+        out.device(dev) = static_cast<T>(1 - epsilon) * in +
+                          static_cast<T>(epsilon / label_dim);
+      }
     }
   }
 };
@@ -54,13 +55,15 @@ class LabelSmoothGradKernel : public framework::OpKernel<T> {
     auto* d_out_t = ctx.Input<framework::Tensor>(framework::GradVarName("Out"));
     auto* d_in_t = ctx.Output<framework::Tensor>(framework::GradVarName("X"));
     d_in_t->mutable_data<T>(ctx.GetPlace());
+    auto d_out_dim = d_out_t->dims()[d_out_t->dims().size() - 1];
+    if (d_out_dim != 0) {
+      auto d_out = framework::EigenVector<T>::Flatten(*d_out_t);
+      auto d_in = framework::EigenVector<T>::Flatten(*d_in_t);
 
-    auto d_out = framework::EigenVector<T>::Flatten(*d_out_t);
-    auto d_in = framework::EigenVector<T>::Flatten(*d_in_t);
-
-    auto epsilon = ctx.Attr<float>("epsilon");
-    auto& dev = *ctx.template device_context<DeviceContext>().eigen_device();
-    d_in.device(dev) = static_cast<T>(1 - epsilon) * d_out;
+      auto epsilon = ctx.Attr<float>("epsilon");
+      auto& dev = *ctx.template device_context<DeviceContext>().eigen_device();
+      d_in.device(dev) = static_cast<T>(1 - epsilon) * d_out;
+    }
   }
 };
 }  // namespace operators
-- 
GitLab