diff --git a/Chapter05/02_Shallow_embeddings.ipynb b/Chapter05/02_Shallow_embeddings.ipynb
index f72c486..ea3d0ad 100644
--- a/Chapter05/02_Shallow_embeddings.ipynb
+++ b/Chapter05/02_Shallow_embeddings.ipynb
@@ -220,9 +220,11 @@
     "        c_tool = 10\n",
     "        \n",
     "        while it < self.max_iter & c_tool > self.tol:\n",
-    "            Y = A*Y_prev\n",
-    "            #force labeled nodes\n",
-    "            Y[labeled_index] = Y0[labeled_index]\n",
+    "            Y = A * Y_prev\n",
+    "            Y = Y / Y.sum(axis=1)  # Normalize rows to sum to 1\n",
+    "            Y[np.isnan(Y)] = 0 # NaN may arise because of all-zeros rows\n",
+    "            \n",
+    "            Y[labeled_index] = Y0[labeled_index] #force labeled nodes\n",
     "            \n",
     "            it +=1\n",
     "            c_tol = np.sum(np.abs(Y-Y_prev))\n",
@@ -396,6 +398,9 @@
     "        while it < self.max_iter & c_tool > self.tol:\n",
     "            Y = self.alpha*(L*Y_prev)+((1-self.alpha)*Y0)\n",
     "\n",
+    "            Y = Y / Y.sum(axis=1)  # Normalize rows to sum to 1\n",
+    "            Y[np.isnan(Y)] = 0 # NaN may arise because of all-zeros rows\n",
+    "\n",
     "            it +=1\n",
     "            c_tol = np.sum(np.abs(Y-Y_prev))\n",
     "            Y_prev = Y\n",