diff --git a/notebooks/254-llm-chatbot/ov_llm_model.py b/notebooks/254-llm-chatbot/ov_llm_model.py
index 59a53ef4024..e7c0d332c83 100644
--- a/notebooks/254-llm-chatbot/ov_llm_model.py
+++ b/notebooks/254-llm-chatbot/ov_llm_model.py
@@ -301,14 +301,7 @@ def __init__(
             model, config, device, dynamic_shapes, ov_config, model_save_dir, **kwargs
         )
     
-    def _reshape(
-        self,
-        model: "Model",
-        batch_size: int,
-        sequence_length: int,
-        height: int = None,
-        width: int = None,
-    ):
+    def _reshape(self, model: "Model", *args, **kwargs):
         shapes = {}
         for inputs in model.inputs:
             shapes[inputs] = inputs.get_partial_shape()
@@ -368,4 +361,4 @@ def _from_pretrained(
     "baichuan2": OVBAICHUANModel,
     "chatglm3": OVCHATGLMModel,
     "gemma": OVModelForCausalLM
-}
\ No newline at end of file
+}