Fix PR comments

sony · Nov 12, 2023 · 1df1636 · 1df1636
1 parent c8ec6c6
commit 1df1636
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/tutorials/notebooks/example_keras_nanodet_plus.ipynb b/tutorials/notebooks/example_keras_nanodet_plus.ipynb
@@ -266,6 +266,7 @@
    "outputs": [],
    "source": [
     "import model_compression_toolkit as mct\n",
+    "from typing import Iterator, Tuple\n",
     "\n",
     "REPRESENTATIVE_DATASET_FOLDER = '/content/coco/val2017'\n",
     "REPRESENTATIVE_DATASET_ANNOTATION_FILE = '/content/coco/annotations/instances_val2017.json'\n",
@@ -278,9 +279,9 @@
     "                                                batch_size=BATCH_SIZE)\n",
     "\n",
     "# Define representative dataset generator\n",
-    "def get_representative_dataset(n_iter, train_loader):\n",
-    "\n",
-    "    def representative_dataset():\n",
+    "def get_representative_dataset(n_iter: int, dataset_loader: Iterator[Tuple]):\n",
+    "    \n",
+    "    def representative_dataset() -> Iterator[List]:\n",
     "        ds_iter = iter(train_loader)\n",
     "        for _ in range(n_iter):\n",
     "            yield [next(ds_iter)[0]]\n",