diff --git a/demo.ipynb b/demo.ipynb
index c2701bf..1d2f469 100644
--- a/demo.ipynb
+++ b/demo.ipynb
@@ -634,7 +634,7 @@
     "    ).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",
+    "    if hasattr(prompt_ids, \"keys\"):\n",
     "        input_ids = prompt_ids[\"input_ids\"]\n",
     "        attention_mask = prompt_ids.get(\"attention_mask\", torch.ones_like(input_ids))\n",
     "    else:\n",