diff --git a/11_training_deep_neural_networks.ipynb b/11_training_deep_neural_networks.ipynb
index 9c7345c..d031822 100644
--- a/11_training_deep_neural_networks.ipynb
+++ b/11_training_deep_neural_networks.ipynb
@@ -1828,8 +1828,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "#with keras.backend.learning_phase_scope(1):  # TODO: check https://github.com/tensorflow/tensorflow/issues/25754\n",
-    "#    history = model.fit(X_train_scaled, y_train)"
+    "history = model.fit(X_train_scaled, y_train)"
    ]
   },
   {
@@ -1855,8 +1854,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "with keras.backend.learning_phase_scope(1):  # TODO: check https://github.com/tensorflow/tensorflow/issues/25754\n",
-    "    y_probas = np.stack([model.predict(X_test_scaled) for sample in range(100)])\n",
+    "y_probas = np.stack([model(X_test_scaled, training=True)\n",
+    "                     for sample in range(100)])\n",
     "y_proba = y_probas.mean(axis=0)\n",
     "y_std = y_probas.std(axis=0)"
    ]