Fix beam_search InferShape (#25169) (#25216)

* fix beam_search infershape, test=develop

* fix beam search op unittest, test=develop

Author: FrostML

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:

python/paddle/fluid/tests/unittests/test_beam_search_op.py

@@ -38,9 +38,9 @@ def setUp(self):
         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(