diff --git a/04_training_linear_models.ipynb b/04_training_linear_models.ipynb
index 72add07..1b1ed32 100644
--- a/04_training_linear_models.ipynb
+++ b/04_training_linear_models.ipynb
@@ -205,6 +205,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import matplotlib.pyplot as plt\n",
+    "\n",
     "plt.figure(figsize=(6, 4))  # not in the book\n",
     "plt.plot(X_new, y_predict, \"r-\", label=\"Predictions\")\n",
     "plt.plot(X, y, \"b.\")\n",
@@ -1052,14 +1054,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "# not in the book\n",
+    "\n",
     "np.random.seed(42)\n",
     "m = 100\n",
     "X = 6 * np.random.rand(m, 1) - 3\n",
     "y = 0.5 * X ** 2 + X + 2 + np.random.randn(m, 1)\n",
-    "# X_train, X_val, y_train, y_val = train_test_split(\n",
-    "# X[:50], y[:50].ravel(), test_size=0.5, random_state=10)\n",
-    "X_train, y_train = X[:50], y[:50, 0]\n",
-    "X_valid, y_valid = X[50:], y[50:, 0]"
+    "X_train, y_train = X[: m // 2], y[: m // 2, 0]\n",
+    "X_valid, y_valid = X[m // 2 :], y[m // 2 :, 0]"
    ]
   },
   {
@@ -1070,6 +1072,7 @@
    "source": [
     "from copy import deepcopy\n",
     "from sklearn.metrics import mean_squared_error\n",
+    "from sklearn.preprocessing import StandardScaler\n",
     "\n",
     "preprocessing = make_pipeline(PolynomialFeatures(degree=90, include_bias=False),\n",
     "                              StandardScaler())\n",
@@ -1192,7 +1195,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "print(iris.data.head())"
+    "iris.data.head(3)"
    ]
   },
   {
@@ -1201,7 +1204,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "iris.target.head()  # note that the instances are not shuffled"
+    "iris.target.head(3)  # note that the instances are not shuffled"
    ]
   },
   {