From 51ef4f3d8d1d2c529922c2f4c611e7a2edf0ff62 Mon Sep 17 00:00:00 2001
From: Sarvesh <sarvesh@Sarveshs-MacBook-Pro-1187.local>
Date: Tue, 9 Jun 2026 16:57:30 +0530
Subject: [PATCH] Use hasattr keys check to support UserDict-based
 BatchEncoding

---
 demo.ipynb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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",