Merge pull request #16389 from phlrain/pick_topk

Merge pull request #16351 from phlrain/fix_topk_shape_check

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");
-    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;