diff --git a/paddle/fluid/operators/masked_select_op.cc b/paddle/fluid/operators/masked_select_op.cc index 3b44c02757fae9648a7e660a06c03af45d621e02..17bf5df18adc543ea487160a31d05d3c802b95a7 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"); }