Merge pull request #16351 from phlrain/fix_topk_shape_check

Fix topk shape check
revert-16190-refine_parallel_executor
Hongyu Liu 6 years ago committed by GitHub
commit 18a0f6d97a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -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;

Loading…
Cancel
Save