diff --git a/10_neural_nets_with_keras.ipynb b/10_neural_nets_with_keras.ipynb
index f26bebf..4419f6a 100644
--- a/10_neural_nets_with_keras.ipynb
+++ b/10_neural_nets_with_keras.ipynb
@@ -1054,8 +1054,8 @@
    "outputs": [],
    "source": [
     "class WideAndDeepModel(keras.models.Model):\n",
-    "    def __init__(self, units=30, activation=\"relu\"):\n",
-    "        super().__init__()\n",
+    "    def __init__(self, units=30, activation=\"relu\", **kwargs):\n",
+    "        super().__init__(**kwargs)\n",
     "        self.hidden1 = keras.layers.Dense(units, activation=activation)\n",
     "        self.hidden2 = keras.layers.Dense(units, activation=activation)\n",
     "        self.main_output = keras.layers.Dense(1)\n",