mypy

Author: xadupre

CHANGELOGS.rst

@@ -4,6 +4,7 @@ Change Logs
 0.7.2
 +++++
 
+* :pr:`166`: improves handling of StaticCache
 * :pr:`165`: support for task text-to-image
 * :pr:`162`: improves graphs rendering for historical data
 

onnx_diagnostic/tasks/text_generation.py

@@ -188,7 +188,7 @@ def get_inputs(
                     (batch_size, num_key_value_heads, sequence_length2, head_dim)
                 ).to(torch.bool),
                 cache_position=torch.arange(sequence_length2).to(torch.int64),
-                past_key_values=make_cache(
+                past_key_values=make_dynamic_cache(
                     [
                         (
                             torch.randn(
@@ -237,7 +237,7 @@ def get_inputs(
                 position_ids=torch.arange(sequence_length, sequence_length + sequence_length2)
                 .to(torch.int64)
                 .expand((batch_size, -1)),
-                past_key_values=make_cache(
+                past_key_values=make_cache(  # type: ignore[operator]
                     [
                         (
                             torch.randn(