diff --git a/src/finn/custom_op/fpgadataflow/pool_batch.py b/src/finn/custom_op/fpgadataflow/pool_batch.py
index ba8a446f2cf7541c0bd2e1dff731afe2397942ef..92679e306ae268c5c6af4399d6dadf5a85f3806b 100644
--- a/src/finn/custom_op/fpgadataflow/pool_batch.py
+++ b/src/finn/custom_op/fpgadataflow/pool_batch.py
@@ -211,7 +211,7 @@ class Pool_Batch(HLSCustomOp):
         self.code_gen_dict["$DEFINES$"] += ["#define PE {}".format(pe)]
 
         k = self.get_nodeattr("KernelSize")
-        self.code_gen_dict["$DEFINES$"] += ["#define KernelSize {}".format(k)]
+        self.code_gen_dict["$DEFINES$"] += ["#define KernelSize {}".format(k * k)]
 
         odim = self.get_nodeattr("OutImgDim")
         self.code_gen_dict["$DEFINES$"] += ["#define OFMDim {}".format(odim)]