diff --git a/02_end_to_end_machine_learning_project.ipynb b/02_end_to_end_machine_learning_project.ipynb
index 571ecb9..4efc429 100644
--- a/02_end_to_end_machine_learning_project.ipynb
+++ b/02_end_to_end_machine_learning_project.ipynb
@@ -59,13 +59,14 @@
     "# Where to save the figures\n",
     "PROJECT_ROOT_DIR = \".\"\n",
     "CHAPTER_ID = \"end_to_end_project\"\n",
+    "IMAGES_PATH = os.path.join(PROJECT_ROOT_DIR, \"images\", CHAPTER_ID)\n",
     "\n",
-    "def save_fig(fig_id, tight_layout=True):\n",
-    "    path = os.path.join(PROJECT_ROOT_DIR, \"images\", CHAPTER_ID, fig_id + \".png\")\n",
+    "def save_fig(fig_id, tight_layout=True, fig_extension=\"png\", resolution=300):\n",
+    "    path = os.path.join(IMAGES_PATH, fig_id + \".\" + fig_extension)\n",
     "    print(\"Saving figure\", fig_id)\n",
     "    if tight_layout:\n",
     "        plt.tight_layout()\n",
-    "    plt.savefig(path, format='png', dpi=300)"
+    "    plt.savefig(path, format=fig_extension, dpi=resolution)"
    ]
   },
   {
@@ -666,7 +667,7 @@
    },
    "outputs": [],
    "source": [
-    "housing_num = housing.drop(\"ocean_proximity\", axis=1)"
+    "housing_num = housing.select_dtypes(include=[np.number])"
    ]
   },
   {
@@ -774,7 +775,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "housing_cat = housing[\"ocean_proximity\"]\n",
+    "housing_cat = housing.select_dtypes(exclude=[np.number])\n",
     "housing_cat.head(10)"
    ]
   },
@@ -2299,11 +2300,12 @@
    "width": "309px"
   },
   "toc": {
-   "navigate_menu": true,
+   "nav_menu": {},
    "number_sections": true,
    "sideBar": true,
-   "threshold": 6,
+   "skip_h1_title": false,
    "toc_cell": false,
+   "toc_position": {},
    "toc_section_display": "block",
    "toc_window_display": false
   }