diff --git a/plonk/gates/poseidon_mds_gate.go b/plonk/gates/poseidon_mds_gate.go
index f03319d..db23e9a 100644
--- a/plonk/gates/poseidon_mds_gate.go
+++ b/plonk/gates/poseidon_mds_gate.go
@@ -90,7 +90,9 @@ func (g *PoseidonMdsGate) EvalUnfiltered(
 	for i := uint64(0); i < poseidon.SPONGE_WIDTH; i++ {
 		output := vars.GetLocalExtAlgebra(g.WireOutput(i))
 		diff := glApi.SubExtensionAlgebra(output, computed_outputs[i])
-		constraints = append(constraints, diff[0], diff[1])
+		for i := 0; i < gl.D; i++ {
+			constraints = append(constraints, diff[i])
+		}
 	}
 
 	return constraints
diff --git a/plonk/gates/public_input_gate.go b/plonk/gates/public_input_gate.go
index 9fad5e7..caf780e 100644
--- a/plonk/gates/public_input_gate.go
+++ b/plonk/gates/public_input_gate.go
@@ -38,7 +38,7 @@ func (g *PublicInputGate) EvalUnfiltered(
 
 	wires := g.WiresPublicInputsHash()
 	hash_parts := vars.publicInputsHash
-	for i := 0; i < 4; i++ {
+	for i := 0; i < len(wires); i++ {
 		wire := wires[i]
 		hash_part := hash_parts[i]