fix squeeze op shape check; test=develop

e7fb344c · phlrain · b7baeed7 · e7fb344c
显示空白变更内容
内联并排

Showing with 5 addition and 2 deletion

paddle/fluid/operators/top_k_op.cc paddle/fluid/operators/top_k_op.cc +5 -2

未找到文件。
--- a/paddle/fluid/operators/top_k_op.cc
+++ b/paddle/fluid/operators/top_k_op.cc
@@ -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");
+    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;