From 3f32b730b2b4972984e9c54d2e1271567bff2ced Mon Sep 17 00:00:00 2001
From: Linjie Chen <40840292+linjieccc@users.noreply.github.com>
Date: Tue, 10 Aug 2021 10:23:35 +0800
Subject: [PATCH] Fix error of HSigmoidLoss (#34719)

* Fix error of HSigmoidLoss

* update unittest

* update unittest
---
 paddle/fluid/operators/hierarchical_sigmoid_op.cc  | 13 +++++++++++--
 .../fluid/tests/unittests/test_hsigmoid_op.py      | 14 ++++++++++++++
 2 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/paddle/fluid/operators/hierarchical_sigmoid_op.cc b/paddle/fluid/operators/hierarchical_sigmoid_op.cc
index 05d521be5a1..b9fbd18cf14 100644
--- a/paddle/fluid/operators/hierarchical_sigmoid_op.cc
+++ b/paddle/fluid/operators/hierarchical_sigmoid_op.cc
@@ -71,8 +71,17 @@ class HierarchicalSigmoidOp : public framework::OperatorWithKernel {
     if (with_prefetch) {
       OP_INOUT_CHECK(ctx->HasOutput("W_Out"), "Output", "W_Out", "hsigmoid");
     }
-    const int64_t batch_size = ctx->GetInputDim("X")[0];
-    std::vector<int64_t> output_shape({batch_size, 1});
+    const int64_t input_dims = ctx->GetInputDim("X")[0];
+    const int64_t label_dims = ctx->GetInputDim("Label")[0];
+    PADDLE_ENFORCE_EQ(input_dims, label_dims,
+                      platform::errors::InvalidArgument(
+                          "The first dimension of "
+                          "input and label is expected to be the same. "
+                          "But received input's first dimension is %d; "
+                          "label's first dimension is %d.",
+                          input_dims, label_dims));
+
+    std::vector<int64_t> output_shape({input_dims, 1});
     ctx->SetOutputDim("Out", framework::make_ddim(output_shape));
     ctx->ShareLoD("X", /*->*/ "Out");
   }
diff --git a/python/paddle/fluid/tests/unittests/test_hsigmoid_op.py b/python/paddle/fluid/tests/unittests/test_hsigmoid_op.py
index 590c3e061f2..965ae65614a 100644
--- a/python/paddle/fluid/tests/unittests/test_hsigmoid_op.py
+++ b/python/paddle/fluid/tests/unittests/test_hsigmoid_op.py
@@ -575,6 +575,20 @@ class TestHSigmoidLossAPI(unittest.TestCase):
                 weight,
                 path_code=path_code_int32)
 
+        # test paddle.nn.HSigmoidLoss
+        paddle.disable_static(self.place)
+        x_arr = np.array([], dtype=np.float32)
+        x = paddle.to_tensor(np.reshape(x_arr, (100000, 0)))
+        label = paddle.to_tensor(0, dtype='int64')
+        self.assertRaises(ValueError, paddle.nn.HSigmoidLoss, x, label)
+
+        # test paddle.nn.functional.hsigmoid_loss
+        x = paddle.to_tensor(np.reshape(x_arr, (10, 0)), dtype='float32')
+        label = paddle.to_tensor([], dtype='int64')
+        weight = paddle.to_tensor([], dtype='float32')
+        self.assertRaises(ValueError, F.hsigmoid_loss, x, label, 0, weight)
+        paddle.enable_static()
+
         # test paddle.fluid.layers.hsigmoid
         with program_guard(Program()):
             label = fluid.data('label', [4, 1], 'int64')
-- 
GitLab