提交 e7fb344c 编写于 作者: P phlrain

fix squeeze op shape check; test=develop

上级 b7baeed7
......@@ -34,8 +34,11 @@ class TopkOp : public framework::OperatorWithKernel {
PADDLE_ENFORCE_GE(k, 1, "k must >= 1");
PADDLE_ENFORCE_GE(input_dims.size(), 1, "input must have >= 1d shape");
PADDLE_ENFORCE_GE(input_dims[input_dims.size() - 1], k,
"input must have >= k columns");
if (ctx->IsRuntime()) {
PADDLE_ENFORCE_GE(input_dims[input_dims.size() - 1], k,
"input must have >= k columns");
}
framework::DDim dims = input_dims;
dims[dims.size() - 1] = k;
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册