From 9d4722c8312cb7e6f06e124c9d254e307da001c4 Mon Sep 17 00:00:00 2001
From: Zhang Zheng <32410583+ZzSean@users.noreply.github.com>
Date: Wed, 2 Jun 2021 12:50:45 +0800
Subject: [PATCH] fix masked_select infer shape (#33167)

---
 paddle/fluid/operators/masked_select_op.cc | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/paddle/fluid/operators/masked_select_op.cc b/paddle/fluid/operators/masked_select_op.cc
index 3b44c02757fae..17bf5df18adc5 100644
--- a/paddle/fluid/operators/masked_select_op.cc
+++ b/paddle/fluid/operators/masked_select_op.cc
@@ -26,8 +26,9 @@ class MaskedSelectOp : public framework::OperatorWithKernel {
     OP_INOUT_CHECK(ctx->HasInput("X"), "Input", "Input", "MaskedSelect");
     OP_INOUT_CHECK(ctx->HasInput("Mask"), "Input", "Mask", "MaskedSelect");
     OP_INOUT_CHECK(ctx->HasOutput("Y"), "Output", "Out", "MaskedSelect");
-    framework::DDim output_dims(ctx->GetInputDim("X"));
-    ctx->SetOutputDim("Y", output_dims);
+
+    // output will only be a 1-D Tensor
+    ctx->SetOutputDim("Y", framework::make_ddim({-1}));
     ctx->ShareLoD("X", /*->*/ "Y");
   }