diff --git a/src/finn/custom_op/fpgadataflow/labelselect_batch.py b/src/finn/custom_op/fpgadataflow/labelselect_batch.py index 9577d2aee7a8333b3cdff2b690270f358d4a10bb..b54f945d3f09c7fe8f5f2bc69bb95ab4cb87a1a3 100644 --- a/src/finn/custom_op/fpgadataflow/labelselect_batch.py +++ b/src/finn/custom_op/fpgadataflow/labelselect_batch.py @@ -69,6 +69,7 @@ class LabelSelect_Batch(HLSCustomOp): pe = self.get_nodeattr("PE") vecs = list(self.get_nodeattr("numInputVectors")) assert nlabels % pe == 0, "PE must divide Labels" + assert nlabels == pe, "LabelSelect currently fails with folding" folds = int(nlabels / pe) folded_ishape = tuple(vecs + [folds, pe]) return folded_ishape