Softmax in functorch example fixed (pytorch#107988)

The output of softmax was overwritten by the output of fc2 in the following line. So, the output of the softmax is never utilized. Now, the final output of the model includes softmax.

Pull Request resolved: pytorch#107988
Approved by: https://github.com/zou3519

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",