Also adopt changes from openvinotoolkit#672

nikita-savelyevv · Aug 20, 2024 · fa051e4 · fa051e4
1 parent fe9789a
commit fa051e4
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/llm_bench/python/llm_bench_utils/conversion_utils/helpers.py b/llm_bench/python/llm_bench_utils/conversion_utils/helpers.py
@@ -96,8 +96,8 @@ def save_tokenizer(tokenizer, out_dir):
 
 def transform_fn(item, item_name, input_shapes, tokenizer, config, max_tokens=127):
     tokenized_text = tokenizer(item[item_name], return_tensors="np")
-    input_ids = tokenized_text["input_ids"][:max_tokens]
-    attention_mask = tokenized_text["attention_mask"][:max_tokens]
+    input_ids = tokenized_text["input_ids"][:, :max_tokens]
+    attention_mask = tokenized_text["attention_mask"][:, :max_tokens]
 
     inputs = {}
     inputs["input_ids"] = input_ids