modify infershape of yolo_box (#40056)

* modify infershape of yolo_box

modify infershape of yolo_box (#40056)
* modify infershape of yolo_box
ebc6959c · wangxinxin08 · GitHub · 7ef61789 · ebc6959c
隐藏空白更改
内联并排

Showing with 6 addition and 1 deletion

paddle/fluid/operators/detection/yolo_box_op.cc paddle/fluid/operators/detection/yolo_box_op.cc +6 -1

未找到文件。
--- a/paddle/fluid/operators/detection/yolo_box_op.cc
+++ b/paddle/fluid/operators/detection/yolo_box_op.cc
@@ -102,7 +102,12 @@ class YoloBoxOp : public framework::OperatorWithKernel {
                          "But received class_num (%s)",
                          class_num));

-    int box_num = dim_x[2] * dim_x[3] * anchor_num;
+    int box_num;
+    if ((dim_x[2] > 0 && dim_x[3] > 0) || ctx->IsRuntime()) {
+      box_num = dim_x[2] * dim_x[3] * anchor_num;
+    } else {
+      box_num = -1;
+    }
    std::vector<int64_t> dim_boxes({dim_x[0], box_num, 4});
    ctx->SetOutputDim("Boxes", phi::make_ddim(dim_boxes));