From 68e93d8a17e4986bab54219662bdc5057d9cc831 Mon Sep 17 00:00:00 2001
From: liu zhengxi <380185688@qq.com>
Date: Fri, 26 Jun 2020 08:13:05 +0800
Subject: [PATCH] Fix beam_search InferShape (#25169)

* fix beam_search infershape, test=develop

* fix beam search op unittest, test=develop
---
 paddle/fluid/operators/beam_search_op.cc                   | 4 ++++
 python/paddle/fluid/tests/unittests/test_beam_search_op.py | 6 +++---
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/paddle/fluid/operators/beam_search_op.cc b/paddle/fluid/operators/beam_search_op.cc
index c6715227cf..887d28f587 100644
--- a/paddle/fluid/operators/beam_search_op.cc
+++ b/paddle/fluid/operators/beam_search_op.cc
@@ -95,6 +95,10 @@ class BeamSearchOp : public framework::OperatorWithKernel {
          std::vector<std::string>({"selected_ids", "selected_scores"})) {
       OP_INOUT_CHECK(ctx->HasOutput(arg), "Output", arg, "BeamSeach");
     }
+    auto id_dims = ctx->GetInputDim("pre_ids");
+    ctx->SetOutputDim("selected_scores", ctx->GetInputDim("pre_scores"));
+    ctx->SetOutputDim("selected_ids", id_dims);
+    ctx->SetOutputDim("parent_idx", {id_dims[0]});
   }
 
  protected:
diff --git a/python/paddle/fluid/tests/unittests/test_beam_search_op.py b/python/paddle/fluid/tests/unittests/test_beam_search_op.py
index faf085eb6c..346cd1e212 100644
--- a/python/paddle/fluid/tests/unittests/test_beam_search_op.py
+++ b/python/paddle/fluid/tests/unittests/test_beam_search_op.py
@@ -38,9 +38,9 @@ class BeamSearchOpTester(unittest.TestCase):
         self._create_pre_scores()
         self._create_scores()
         self._create_pre_ids()
-        self.scope.var('selected_ids')
-        self.scope.var('selected_scores')
-        self.scope.var('parent_idx')
+        self.scope.var('selected_ids').get_tensor()
+        self.scope.var('selected_scores').get_tensor()
+        self.scope.var('parent_idx').get_tensor()
 
     def test_run(self):
         op = Operator(
-- 
GitLab