fix infer shapes of pool_with_index (#40139)

* dbg pool infer shapes * dbg * fix format
PaddlePaddle · Mar 7, 2022 · 0fb6bca · 0fb6bca
1 parent 55a3bfb
commit 0fb6bca
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/paddle/fluid/operators/pool_with_index_op.cc b/paddle/fluid/operators/pool_with_index_op.cc
@@ -81,8 +81,12 @@ class MaxPoolWithIndexOp : public framework::OperatorWithKernel {
       output_shape.insert(output_shape.end(), ksize.begin(), ksize.end());
     } else {
       for (size_t i = 0; i < ksize.size(); ++i) {
-        output_shape.push_back(MaxPoolOutputSize(in_x_dims[i + 2], ksize[i],
-                                                 paddings[i], strides[i]));
+        if ((!ctx->IsRuntime()) && (in_x_dims[i + 2] < 0)) {
+          output_shape.push_back(in_x_dims[i + 2]);
+        } else {
+          output_shape.push_back(MaxPoolOutputSize(in_x_dims[i + 2], ksize[i],
+                                                   paddings[i], strides[i]));
+        }
       }
     }
     ctx->SetOutputDim("Out", phi::make_ddim(output_shape));