diff --git a/paddle/operators/dropout_op.cc b/paddle/operators/dropout_op.cc index c4bee44e3e..82e2140808 100644 --- a/paddle/operators/dropout_op.cc +++ b/paddle/operators/dropout_op.cc @@ -25,8 +25,6 @@ class DropoutOp : public framework::OperatorWithKernel { void InferShape(framework::InferShapeContext* ctx) const override { PADDLE_ENFORCE(ctx->HasInput("X"), "Input(X) must not be null."); - PADDLE_ENFORCE_GE(ctx->Attrs().Get("dropout_prob"), 0); - PADDLE_ENFORCE_LE(ctx->Attrs().Get("dropout_prob"), 1); auto x_dims = ctx->GetInputDim("X"); ctx->SetOutputDim("Out", x_dims); @@ -47,7 +45,11 @@ class DropoutOpMaker : public framework::OpProtoAndCheckerMaker { AddOutput("Mask", "The random sampled dropout mask.").AsIntermediate(); AddAttr("dropout_prob", "Probability of setting units to zero.") - .SetDefault(.5f); + .SetDefault(.5f) + .AddCustomChecker([](const float& drop_p) { + PADDLE_ENFORCE(drop_p > 0.0f && drop_p < 1.0f, + "'dropout_prob' must be between 0 and 1."); + }); AddAttr("is_test", "True if in test phase.").SetDefault(false); AddAttr("seed", "Dropout random seed.").SetDefault(0); @@ -78,8 +80,6 @@ class DropoutOpGrad : public framework::OperatorWithKernel { PADDLE_ENFORCE(ctx->HasInput(framework::GradVarName("Out")), "Input(Out@GRAD) must not be null."); - PADDLE_ENFORCE_GE(ctx->Attrs().Get("dropout_prob"), 0); - PADDLE_ENFORCE_LE(ctx->Attrs().Get("dropout_prob"), 1); auto x_dims = ctx->GetInputDim("X"); auto out_dims = ctx->GetInputDim(framework::GradVarName("Out")); PADDLE_ENFORCE_EQ(x_dims, out_dims,