diff --git a/demo.ipynb b/demo.ipynb
index d0a1b10..c2701bf 100644
--- a/demo.ipynb
+++ b/demo.ipynb
@@ -633,9 +633,18 @@
     "        return_tensors=\"pt\"\n",
     "    ).to(device)\n",
     "\n",
+    "    # Handle cases where apply_chat_template returns a BatchEncoding dict instead of a raw Tensor\n",
+    "    if isinstance(prompt_ids, dict):\n",
+    "        input_ids = prompt_ids[\"input_ids\"]\n",
+    "        attention_mask = prompt_ids.get(\"attention_mask\", torch.ones_like(input_ids))\n",
+    "    else:\n",
+    "        input_ids = prompt_ids\n",
+    "        attention_mask = torch.ones_like(input_ids)\n",
+    "\n",
+    "\n",
     "    generated_ids = model.base_model.generate(\n",
-    "        input_ids=prompt_ids,\n",
-    "        attention_mask=torch.ones_like(prompt_ids),\n",
+    "        input_ids=input_ids,\n",
+    "        attention_mask=attention_mask,\n",
     "        max_new_tokens=train_args.generation_max_new_tokens,\n",
     "        do_sample=False,\n",
     "        pad_token_id=tokenizer.pad_token_id,\n",
@@ -643,7 +652,7 @@
     "    )\n",
     "\n",
     "    video2lora_prediction = tokenizer.decode(\n",
-    "        generated_ids[0][prompt_ids.shape[1]:],\n",
+    "        generated_ids[0][input_ids.shape[1]:],\n",
     "        skip_special_tokens=True\n",
     "    ).strip()\n",
     "\n",