diff --git a/llm_bench/python/utils/conversion_utils/helpers.py b/llm_bench/python/utils/conversion_utils/helpers.py
index 2c7508b6d4..98a39508c3 100644
--- a/llm_bench/python/utils/conversion_utils/helpers.py
+++ b/llm_bench/python/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