diff --git a/tests/transformation/test_move_flatten_past_affine.py b/tests/transformation/test_move_flatten_past_affine.py
index e5588965ce1c2d006529b0f8895c489e7a38b6c7..b2d5e51613d41f3f2db3dabcef7b982ec2816b19 100644
--- a/tests/transformation/test_move_flatten_past_affine.py
+++ b/tests/transformation/test_move_flatten_past_affine.py
@@ -93,7 +93,7 @@ def test_move_flatten_past_affine(data_layout, batch_size):
 
     # compare execution before and after transformation
     inp_values = gen_finn_dt_tensor(DataType.INT2, ishape)
-    idict = {"inp": inp_values}
+    idict = {model.graph.input[0].name: inp_values}
     model_transformed = model.transform(MoveFlattenPastAffine())
     assert oxe.compare_execution(model, model_transformed, idict)