From 0fe1bf6e251ea4c2c1fbe32f2330edd9c1552eff Mon Sep 17 00:00:00 2001
From: Marcos Luciano <marcoslucianops@gmail.com>
Date: Fri, 15 Jul 2022 03:57:47 -0300
Subject: [PATCH] Fix MaxPool layer

---
 nvdsinfer_custom_impl_Yolo/layers/maxpool_layer.cpp | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/nvdsinfer_custom_impl_Yolo/layers/maxpool_layer.cpp b/nvdsinfer_custom_impl_Yolo/layers/maxpool_layer.cpp
index c27308a..51caab9 100644
--- a/nvdsinfer_custom_impl_Yolo/layers/maxpool_layer.cpp
+++ b/nvdsinfer_custom_impl_Yolo/layers/maxpool_layer.cpp
@@ -23,9 +23,11 @@ nvinfer1::ILayer* maxpoolLayer(
     assert(pool);
     std::string maxpoolLayerName = "maxpool_" + std::to_string(layerIdx);
     pool->setStrideNd(nvinfer1::Dims{2, {stride, stride}});
-    if (stride == 1)
+    pool->setPaddingNd(nvinfer1::Dims{2, {(size - 1) / 2, (size - 1) / 2}});
+    if (size == 2 && stride == 1)
     {
-        pool->setPaddingNd(nvinfer1::Dims{2, {size / 2, size / 2}});
+        pool->setPrePadding(nvinfer1::Dims{2, {0, 0}});
+        pool->setPostPadding(nvinfer1::Dims{2, {1, 1}});
     }
     pool->setName(maxpoolLayerName.c_str());