From c177040c77cb561922bab8e8abf74f6804e7678a Mon Sep 17 00:00:00 2001 From: Pawel Raasz Date: Thu, 22 Aug 2024 12:40:01 +0200 Subject: [PATCH] Correct get tensor name for stateful key, values (#874) --- optimum/exporters/openvino/stateful.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/optimum/exporters/openvino/stateful.py b/optimum/exporters/openvino/stateful.py index 8ca42b67aa..d8132107a1 100644 --- a/optimum/exporters/openvino/stateful.py +++ b/optimum/exporters/openvino/stateful.py @@ -200,10 +200,10 @@ def patch_stateful(config: PretrainedConfig, ov_model: ov.Model): """ key_value_input_names = [ - key.get_any_name() for key in ov_model.inputs if any("key_values" in key_name for key_name in key.get_names()) + key_name for key in ov_model.inputs for key_name in key.get_names() if "key_values" in key_name ] key_value_output_names = [ - key.get_any_name() for key in ov_model.outputs if any("present" in key_name for key_name in key.get_names()) + key_name for key in ov_model.outputs for key_name in key.get_names() if "present" in key_name ] not_kv_inputs = [ input for input in ov_model.inputs if not any(name in key_value_input_names for name in input.get_names())