diff --git a/src/finn/custom_op/fpgadataflow/checksum.py b/src/finn/custom_op/fpgadataflow/checksum.py
index 21a09a1c5e3ea6b3d6076798351ef31e4c76c062..c927c07df21faf40ccbf9ddbe47e3f2f2ca61c89 100644
--- a/src/finn/custom_op/fpgadataflow/checksum.py
+++ b/src/finn/custom_op/fpgadataflow/checksum.py
@@ -98,7 +98,7 @@ class CheckSum(HLSCustomOp):
     def get_folded_input_shape(self, ind=0):
         return self.get_nodeattr("folded_shape")
 
-    def get_folded_output_shape(self):
+    def get_folded_output_shape(self, ind=0):
         return self.get_nodeattr("folded_shape")
 
     def get_normal_input_shape(self, ind=0):