diff --git a/12_custom_models_and_training_with_tensorflow.ipynb b/12_custom_models_and_training_with_tensorflow.ipynb
index 84fb79e..f7dbab4 100644
--- a/12_custom_models_and_training_with_tensorflow.ipynb
+++ b/12_custom_models_and_training_with_tensorflow.ipynb
@@ -3365,7 +3365,10 @@
    "source": [
     "# extra code – the model can be loaded and you can continue training or use it\n",
     "#              to make predictions\n",
-    "model = tf.keras.models.load_model(\"my_custom_model.keras\")\n",
+    "model = tf.keras.models.load_model(\n",
+    "    \"my_custom_model.keras\",\n",
+    "    custom_objects={\"ResidualRegressor\": ResidualRegressor}\n",
+    ")\n",
     "history = model.fit(X_train_scaled, y_train, epochs=2)\n",
     "model.predict(X_test_scaled[:3])"
    ]