fixes karpathy#59

Fixes karpathy#59: Corrected the shape of `dbnmeani`, which should be [1, 64] instead of [64].
conscell · Dec 27, 2024 · 9ae6af2 · 9ae6af2
1 parent 73c3fcc
commit 9ae6af2
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/lectures/makemore/makemore_part4_backprop.ipynb b/lectures/makemore/makemore_part4_backprop.ipynb
@@ -323,7 +323,7 @@
     "dbndiff2 = (1.0/(n-1))*torch.ones_like(bndiff2) * dbnvar\n",
     "dbndiff += (2*bndiff) * dbndiff2\n",
     "dhprebn = dbndiff.clone()\n",
-    "dbnmeani = (-dbndiff).sum(0)\n",
+    "dbnmeani = (-dbndiff).sum(0, keepdim=True)\n",
     "dhprebn += 1.0/n * (torch.ones_like(hprebn) * dbnmeani)\n",
     "dembcat = dhprebn @ W1.T\n",
     "dW1 = embcat.T @ dhprebn\n",
@@ -598,7 +598,7 @@
     "# dbndiff2 = (1.0/(n-1))*torch.ones_like(bndiff2) * dbnvar\n",
     "# dbndiff += (2*bndiff) * dbndiff2\n",
     "# dhprebn = dbndiff.clone()\n",
-    "# dbnmeani = (-dbndiff).sum(0)\n",
+    "# dbnmeani = (-dbndiff).sum(0, keepdim=True)\n",
     "# dhprebn += 1.0/n * (torch.ones_like(hprebn) * dbnmeani)\n",
     "\n",
     "# calculate dhprebn given dhpreact (i.e. backprop through the batchnorm)\n",