fix

Author: xadupre

_unittests/ut_tasks/test_tasks_text_generation.py

@@ -54,6 +54,7 @@ def test_text_generation_tiny_llm(self):
         data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True)
         self.assertEqual(data["task"], "text-generation")
         model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
+        inputs_copied = torch_deepcopy(inputs)
         expected = model(**torch_deepcopy(inputs))
         model(**data["inputs2"])
         fake = make_fake_with_dynamic_dimensions(inputs, dynamic_shapes=ds)[0]
@@ -62,7 +63,7 @@ def test_text_generation_tiny_llm(self):
                 model, (), kwargs=fake, dynamic_shapes=use_dyn_not_str(ds), strict=False
             )
             # print(ep)
-        got = ep.module()(**inputs)
+        got = ep.module()(**inputs_copied)
         self.assertEqualAny(expected.past_key_values, got.past_key_values)
         self.assertEqualArray(expected.logits, got.logits)
 

onnx_diagnostic/helpers/cache_helper.py

@@ -182,6 +182,11 @@ def make_dynamic_cache(
                 layer.device = k.device
                 layer.keys = k
                 layer.values = v
+                layer.is_initialized = True
+            assert not hasattr(cache, "layers") or len(key_value_pairs) == len(cache.layers), (
+                f"Unexpected number of layers in the cache ({len(cache.layers)}), "
+                f"{len(key_value_pairs)} expected."
+            )
             return finalize_cache(cache)
 
         cache = transformers.cache_utils.DynamicCache(key_value_pairs)