diff --git a/ProgrammingAssignment_2/model.ipynb b/ProgrammingAssignment_2/model.ipynb
index 413c74b98a01ce2e8c248a6ed50d7af9a660efd2..c37e9bd290d748385d0517d4a5f970e422c765af 100644
--- a/ProgrammingAssignment_2/model.ipynb
+++ b/ProgrammingAssignment_2/model.ipynb
@@ -86,12 +86,12 @@
    "source": [
     "# TODO: Programming Assignment 2\n",
     "\n",
-    "def kfold(indices, k):\n",
+    "def kfold(size, k):\n",
     "\n",
     "    '''\n",
     "    Args:\n",
-    "        indices: ndarray\n",
-    "            1D array with integer entries containing indices\n",
+    "        size: int\n",
+    "            number of examples in the dataset that you want to split into k\n",
     "        k: int \n",
     "            Number of desired splits in data.(Assume test set is already separated.)\n",
     "        Returns:\n",