diff --git a/examples/sentiment/notebooks/gpt2-sentiment.ipynb b/examples/sentiment/notebooks/gpt2-sentiment.ipynb
index 0b56febbdc..b164fbd5a7 100644
--- a/examples/sentiment/notebooks/gpt2-sentiment.ipynb
+++ b/examples/sentiment/notebooks/gpt2-sentiment.ipynb
@@ -62,6 +62,7 @@
    "source": [
     "import torch\n",
     "from tqdm import tqdm\n",
+    "import pandas as pd\n",
     "tqdm.pandas()\n",
     "\n",
     "from transformers import pipeline, AutoTokenizer\n",
@@ -685,7 +686,7 @@
     "dataset.set_format(\"pandas\")\n",
     "df_batch = dataset[:].sample(bs)\n",
     "game_data['query'] = df_batch['query'].tolist()\n",
-    "query_tensors = df_batch['tokens'].tolist()\n",
+    "query_tensors = df_batch['input_ids'].tolist()\n",
     "\n",
     "response_tensors_ref, response_tensors = [], []\n",
     "\n",