Skip to content

Commit dc02405

Browse files
committed
bring back inputs2
1 parent 1f4ca3a commit dc02405

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

onnx_diagnostic/tasks/text_generation.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -262,8 +262,9 @@ def get_inputs(
262262
# which it becomes an empty tensor
263263
res = dict(inputs=inputs, dynamic_shapes=shapes)
264264
if add_second_input:
265+
# TODO(titaiwang): Make input key more informative
265266
# prompt processing (prefill) testing
266-
res["prompt_processing"] = get_inputs(
267+
res["inputs2"] = get_inputs(
267268
model=model,
268269
config=config,
269270
dummy_max_token_id=dummy_max_token_id,

0 commit comments

Comments
 (0)