fix patched lazy_initialization for transformers>=5 (#376)

* fix patched lazy_initialization for transformers>=5

* doc

Author: xadupre

CHANGELOGS.rst

@@ -4,6 +4,7 @@ Change Logs
 0.8.8
 +++++
 
+* :pr:`376`: fix patched lazy_initialization for transformers>=5
 * :pr:`372`: fix patch on rotary embedding
 * :pr:`371`: fix make_fake_with_dynamic_dimensions
 

onnx_diagnostic/torch_export_patches/patches/_patch_transformers_dynamic_cache.py

@@ -22,7 +22,9 @@ class patched_DynamicLayer:
         _PATCHES_ = ["lazy_initialization"]
         _PATCHED_CLASS_ = DynamicLayer
 
-        def lazy_initialization(self, key_states: torch.Tensor):
+        def lazy_initialization(
+            self, key_states: torch.Tensor, value_states: torch.Tensor = None
+        ):
             self.dtype, self.device = key_states.dtype, key_states.device
             assert (
                 hasattr(key_states, "shape") and key_states is not None