diff --git a/paddle/fluid/operators/reshape_op.h b/paddle/fluid/operators/reshape_op.h index ccd7063fe69e0f21b4d2a821bb70902b39c9b9de..fc693317d19254aeb33a0f67df509790ece31e8c 100644 --- a/paddle/fluid/operators/reshape_op.h +++ b/paddle/fluid/operators/reshape_op.h @@ -92,14 +92,16 @@ class ReshapeOp : public framework::OperatorWithKernel { } if (unk_dim_idx != -1) { - output_shape[unk_dim_idx] = -in_size / capacity; - // in_size < 0 and is un-determinate in compile time, skip the check, - // for example, in_dims = [-1, 8, 1, 1], shape = [-1, 3, 8], - // capacity = -24, in_size = -8, output_shape[0] = 0 - // the following check will fail. if (in_size > 0) { + // in_size < 0 and is un-determinate in compile time, skip the check, + // for example, in_dims = [-1, 8, 1, 1], shape = [-1, 3, 8], + // capacity = -24, in_size = -8, output_shape[0] = 0 + // the following check will fail. + output_shape[unk_dim_idx] = -in_size / capacity; PADDLE_ENFORCE_EQ(output_shape[unk_dim_idx] * capacity, -in_size, "Invalid shape is given."); + } else { + output_shape[unk_dim_idx] = -1; } } else { PADDLE_ENFORCE_EQ(capacity, in_size, "Invalid shape is given.");