diff --git a/functorch/notebooks/per_sample_grads.ipynb b/functorch/notebooks/per_sample_grads.ipynb
index 5f7ad23880b5e8..a34c80d07ac4ac 100644
--- a/functorch/notebooks/per_sample_grads.ipynb
+++ b/functorch/notebooks/per_sample_grads.ipynb
@@ -60,7 +60,7 @@
         "        x = self.fc1(x)\n",
         "        x = F.relu(x)\n",
         "        x = self.fc2(x)\n",
-        "        output = F.log_softmax(x, dim=1)\n",
+        "        x = F.log_softmax(x, dim=1)\n",
         "        output = x\n",
         "        return output\n",
         "\n",