diff --git a/project/starter.ipynb b/project/starter.ipynb
index c782396b9..91d41f2eb 100644
--- a/project/starter.ipynb
+++ b/project/starter.ipynb
@@ -196,6 +196,15 @@
     "))"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%matplotlib inline"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -751,10 +760,10 @@
     "\n",
     "```python\n",
     "{\n",
-    "    \"inferences\": [], # Output of predictor.predict\n",
-    "    \"s3_key\": \"\", # Source data S3 key\n",
     "    \"s3_bucket\": \"\", # Source data S3 bucket\n",
-    "    \"image_data\": \"\"  # base64 encoded string containing the image data\n",
+    "    \"s3_key\": \"\", # Source data S3 key\n",
+    "    \"image_data\": \"\",  # base64 encoded string containing the image data\n",
+    "    \"inferences\": [] # Output of predictor.predict\n",
     "}\n",
     "```\n",
     "\n",
@@ -762,9 +771,9 @@
     "\n",
     "```python\n",
     "{\n",
-    "  \"image_data\": \"\",\n",
     "  \"s3_bucket\": MY_BUCKET_NAME, # Fill in with your bucket\n",
-    "  \"s3_key\": \"test/bicycle_s_000513.png\"\n",
+    "  \"s3_key\": \"test/bicycle_s_000513.png\",\n",
+    "  \"image_data\": \"\"\n",
     "}\n",
     "```\n",
     "\n",
@@ -798,9 +807,9 @@
     "    return {\n",
     "        'statusCode': 200,\n",
     "        'body': {\n",
-    "            \"image_data\": image_data,\n",
     "            \"s3_bucket\": bucket,\n",
     "            \"s3_key\": key,\n",
+    "            \"image_data\": image_data,\n",
     "            \"inferences\": []\n",
     "        }\n",
     "    }\n",
@@ -836,10 +845,10 @@
     "    inferences = ## TODO: fill in\n",
     "    \n",
     "    # We return the data back to the Step Function    \n",
-    "    event[\"inferences\"] = inferences.decode('utf-8')\n",
+    "    event[\"body\"][\"inferences\"] = inferences.decode('utf-8')\n",
     "    return {\n",
     "        'statusCode': 200,\n",
-    "        'body': json.dumps(event)\n",
+    "        'body': json.dumps(event[\"body\"])\n",
     "    }\n",
     "```\n",
     "\n",
@@ -869,7 +878,7 @@
     "\n",
     "    return {\n",
     "        'statusCode': 200,\n",
-    "        'body': json.dumps(event)\n",
+    "        'body': json.dumps(event[\"body\"])\n",
     "    }\n",
     "```\n",
     "Once you have tested the lambda functions, save the code for each lambda function in a python script called 'lambda.py'.\n",
@@ -988,12 +997,15 @@
     "import os\n",
     "\n",
     "# List the file names we downloaded\n",
-    "file_handles = os.listdir(\"./captured_data\")\n",
+    "file_handles = []\n",
+    "for root, subFolders, files in os.walk(\"./captured_data\"):\n",
+    "    if files:\n",
+    "        file_handles += [root + '/' + file for file in files]\n",
     "\n",
     "# Dump all the data into an array\n",
     "json_data = []\n",
     "for jsonl in file_handles:\n",
-    "    with jsonlines.open(f\"./captured_data/{jsonl}\") as f:\n",
+    "    with jsonlines.open(jsonl) as f:\n",
     "        json_data.append(f.read())"
    ]
   },
@@ -1084,9 +1096,9 @@
  "metadata": {
   "instance_type": "ml.t3.medium",
   "kernelspec": {
-   "display_name": "Python 3 (Data Science)",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
-   "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:us-east-1:081325390199:image/datascience-1.0"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
@@ -1098,7 +1110,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.10"
+   "version": "3.7.6"
   }
  },
  "nbformat": 4,