fix cache dynamic dimension in image-text-to-txt

xadupre · xadupre · commit caab6da05245 · 2026-02-23T11:13:51.000+01:00
diff --git a/onnx_diagnostic/tasks/image_text_to_text.py b/onnx_diagnostic/tasks/image_text_to_text.py
@@ -156,7 +156,7 @@ def _get_inputs_gemma3(
         },
         "position_ids": {0: batch, 1: seq_length},
         "cache_position": {0: seq_length},
-        "past_key_values": [{0: batch} for _ in range(num_hidden_layers * 2)],
+        "past_key_values": [{0: batch, 2: seq_length} for _ in range(num_hidden_layers * 2)],
         "pixel_values": {0: batch},
         "use_cache": None,
     }
@@ -280,7 +280,7 @@ def get_inputs_default(
         "past_key_values": list(
             itertools.chain.from_iterable(
                 zip(
-                    [{0: batch} for _ in range(num_hidden_layers)],
+                    [{0: batch, 2: cache_length} for _ in range(num_hidden_layers)],
                     [{0: batch, 2: cache_length} for _ in range(num_hidden_layers)],
                 )
             )

Original file line number	Diff line number	Diff line change
`@@ -156,7 +156,7 @@ def _get_inputs_gemma3(`
`156`	`156`	`},`
`157`	`157`	`"position_ids": {0: batch, 1: seq_length},`
`158`	`158`	`"cache_position": {0: seq_length},`
`159`		`- "past_key_values": [{0: batch} for _ in range(num_hidden_layers * 2)],`
	`159`	`+ "past_key_values": [{0: batch, 2: seq_length} for _ in range(num_hidden_layers * 2)],`
`160`	`160`	`"pixel_values": {0: batch},`
`161`	`161`	`"use_cache": None,`
`162`	`162`	`}`
`@@ -280,7 +280,7 @@ def get_inputs_default(`
`280`	`280`	`"past_key_values": list(`
`281`	`281`	`itertools.chain.from_iterable(`
`282`	`282`	`zip(`
`283`		`- [{0: batch} for _ in range(num_hidden_layers)],`
	`283`	`+ [{0: batch, 2: cache_length} for _ in range(num_hidden_layers)],`
`284`	`284`	`[{0: batch, 2: cache_length} for _ in range(num_hidden_layers)],`
`285`	`285`	`)`
`286`	`286`	`)`