diff --git a/02_end_to_end_machine_learning_project.ipynb b/02_end_to_end_machine_learning_project.ipynb
index af1ceba..c63e8ef 100644
--- a/02_end_to_end_machine_learning_project.ipynb
+++ b/02_end_to_end_machine_learning_project.ipynb
@@ -689,7 +689,7 @@
    },
    "outputs": [],
    "source": [
-    "housing_num = housing.drop(\"ocean_proximity\", axis=1)"
+    "housing_num = housing.select_dtypes(include=[np.number])"
    ]
   },
   {
@@ -797,7 +797,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "housing_cat = housing[\"ocean_proximity\"]\n",
+    "housing_cat = housing.select_dtypes(exclude=[np.number])\n",
     "housing_cat.head(10)"
    ]
   },