Fix static cache

Author: xadupre

_unittests/ut_torch_models/test_tiny_llms.py

@@ -56,7 +56,7 @@ def test_tiny_llm_export_static(self):
         self.assertEqual(
             {"attention_mask", "past_key_values", "input_ids", "cache_position"}, set(inputs)
         )
-        with torch_export_patches(patch_transformers=True, stop_if_static=1):
+        with torch_export_patches(patch_transformers=True, stop_if_static=0):
             ep = torch.export.export(
                 model,
                 (),

onnx_diagnostic/tasks/text_generation.py

@@ -176,8 +176,10 @@ def get_inputs(
                 "attention_mask": {0: batch, 2: "seq"},
                 "cache_position": {0: "seq"},
                 "past_key_values": [
-                    [{0: batch, 2: cache_length} 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)],
+                    # [{0: batch, 2: cache_length} for _ in range(num_hidden_layers)],
+                    [{0: batch} for _ in range(num_hidden_layers)],
+                    [{0: batch} for _ in range(num_hidden_layers)],
                 ],
             }
             inputs = dict(