Fix batch switch (#1182) (#1185)

OV didn't allow increasing shape of an already allocated tensor. Ticket 157039 Co-authored-by: Ilya Lavrenov <[email protected]>
openvinotoolkit · Nov 11, 2024 · 790ca60 · 790ca60
1 parent 4017c8f
commit 790ca60
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/src/cpp/src/lm_encoding.cpp b/src/cpp/src/lm_encoding.cpp
@@ -86,7 +86,6 @@ std::pair<EncodedResults, int32_t> get_lm_encoded_results(
     if (position_ids.has_value())
         m_llm.set_tensor("position_ids", *position_ids);
 
-    m_llm.get_tensor("beam_idx").set_shape({ batch_size });
     ov::Tensor beam_idx = ov::Tensor(ov::element::i32, {batch_size});
     auto beam_data = beam_idx.data<int32_t>();
     if (selected_beam_idx.has_value())

diff --git a/tests/python_tests/test_generate_api.py b/tests/python_tests/test_generate_api.py
@@ -840,3 +840,11 @@ def test_perf_metrics(model_descr, generation_config, prompt):
     assert len(raw_metrics.m_times_to_first_token) > 0
     assert len(raw_metrics.m_batch_sizes) > 0
     assert len(raw_metrics.m_durations) > 0
+
+
+@pytest.mark.precommit
+@pytest.mark.nightly
+def test_batch_switch():
+    pipe = read_model(('katuni4ka/tiny-random-phi3', Path('tiny-random-phi3')))[4]
+    pipe.generate(["a"], max_new_tokens=2)
+    pipe.generate(["1", "2"], max_new_tokens=2)