diff --git a/notebooks/fine_tune_with_hugging_face.ipynb b/notebooks/fine_tune_with_hugging_face.ipynb
index 8fe9b30..d5ce5b4 100644
--- a/notebooks/fine_tune_with_hugging_face.ipynb
+++ b/notebooks/fine_tune_with_hugging_face.ipynb
@@ -1027,8 +1027,8 @@
       "outputs": [],
       "source": [
         "pt_outputs = pt_pipe(\n",
-        "    text=test_data[\"messages\"],\n",
-        "    images=test_data[\"image\"],\n",
+        "    text=[txt for txt in test_data[\"messages\"]],\n",
+        "    images=[img for img in test_data[\"image\"]],\n",
         "    max_new_tokens=40,\n",
         "    batch_size=64,\n",
         "    return_full_text=False,\n",
@@ -1124,8 +1124,8 @@
       "outputs": [],
       "source": [
         "ft_outputs = ft_pipe(\n",
-        "    text=test_data[\"messages\"],\n",
-        "    images=test_data[\"image\"],\n",
+        "    text=[txt for txt in test_data[\"messages\"]],\n",
+        "    images=[img for img in test_data[\"image\"]],\n",
         "    max_new_tokens=20,\n",
         "    batch_size=64,\n",
         "    return_full_text=False,\n",