diff --git a/AnnotatedTransformer.ipynb b/AnnotatedTransformer.ipynb
index 0f7da7d..8907584 100644
--- a/AnnotatedTransformer.ipynb
+++ b/AnnotatedTransformer.ipynb
@@ -1514,37 +1514,13 @@
      "output_type": "stream",
      "text": [
       "Example Untrained Model Prediction: tensor([[ 0, 10,  0, 10,  0,  0,  0,  0,  0, 10]])\n",
-      "Example Untrained Model Prediction: tensor([[ 0,  8,  1, 10,  0,  8,  1, 10,  0,  8]])\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
+      "Example Untrained Model Prediction: tensor([[ 0,  8,  1, 10,  0,  8,  1, 10,  0,  8]])\n",
       "Example Untrained Model Prediction: tensor([[ 0,  9,  0, 10,  4,  5,  3,  2,  4,  3]])\n",
-      "Example Untrained Model Prediction: tensor([[0, 5, 5, 5, 5, 5, 5, 5, 5, 5]])\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
+      "Example Untrained Model Prediction: tensor([[0, 5, 5, 5, 5, 5, 5, 5, 5, 5]])\n",
       "Example Untrained Model Prediction: tensor([[0, 2, 8, 3, 8, 5, 0, 4, 0, 4]])\n",
-      "Example Untrained Model Prediction: tensor([[ 0, 10,  3, 10,  2,  9,  0,  3, 10,  3]])\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
+      "Example Untrained Model Prediction: tensor([[ 0, 10,  3, 10,  2,  9,  0,  3, 10,  3]])\n",
       "Example Untrained Model Prediction: tensor([[0, 3, 3, 3, 3, 3, 3, 3, 3, 3]])\n",
-      "Example Untrained Model Prediction: tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
+      "Example Untrained Model Prediction: tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])\n",
       "Example Untrained Model Prediction: tensor([[0, 3, 2, 2, 2, 4, 0, 3, 1, 3]])\n",
       "Example Untrained Model Prediction: tensor([[0, 6, 6, 6, 6, 6, 6, 6, 6, 6]])\n"
      ]
@@ -2209,7 +2185,9 @@
     "            [0, 0.2, 0.7, 0.1, 0],\n",
     "        ]\n",
     "    )\n",
-    "    crit(x=predict.log(), target=torch.LongTensor([2, 1, 0, 3, 3]))\n",
+    "    crit(x=log_softmax(predict), target=torch.LongTensor([2, 1, 0, 3, 3]))\n",
+    "    \n",
+    "    \n",
     "    LS_data = pd.concat(\n",
     "        [\n",
     "            pd.DataFrame(\n",
@@ -2337,8 +2315,8 @@
     "def loss(x, crit):\n",
     "    d = x + 3 * 1\n",
     "    predict = torch.FloatTensor([[0, x / d, 1 / d, 1 / d, 1 / d]])\n",
-    "    return crit(predict.log(), torch.LongTensor([1])).data\n",
-    "\n",
+    "    return crit( log_softmax(predict), torch.LongTensor([1])).data\n",
+    "    \n",
     "\n",
     "def penalization_visualization():\n",
     "    crit = LabelSmoothing(5, 0, 0.1)\n",
@@ -3566,38 +3544,14 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Preparing Data ...\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Loading Trained Model ...\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
+      "Preparing Data ...\n",
+      "Loading Trained Model ...\n",
       "Checking Model Outputs:\n",
       "\n",
       "Example 0 ========\n",
-      "\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
+      "\n",
       "Source Text (Input)        : <s> Mehrere Kinder heben die Hände , während sie auf einem bunten Teppich in einem Klassenzimmer sitzen . </s>\n",
-      "Target Text (Ground Truth) : <s> Several children are raising their hands while sitting on a colorful rug in a classroom . </s>\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
+      "Target Text (Ground Truth) : <s> Several children are raising their hands while sitting on a colorful rug in a classroom . </s>\n",
       "Model Output               : <s> A group of children are in their hands while sitting on a colorful carpet . </s>\n"
      ]
     },
@@ -3716,38 +3670,14 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Preparing Data ...\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Loading Trained Model ...\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
+      "Preparing Data ...\n",
+      "Loading Trained Model ...\n",
       "Checking Model Outputs:\n",
       "\n",
       "Example 0 ========\n",
-      "\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
+      "\n",
       "Source Text (Input)        : <s> Drei Menschen wandern auf einem stark verschneiten Weg . </s>\n",
-      "Target Text (Ground Truth) : <s> A <unk> of people are hiking throughout a heavily snowed path . </s>\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
+      "Target Text (Ground Truth) : <s> A <unk> of people are hiking throughout a heavily snowed path . </s>\n",
       "Model Output               : <s> Three people hiking on a busy <unk> . </s>\n"
      ]
     },
@@ -3869,38 +3799,14 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Preparing Data ...\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Loading Trained Model ...\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
+      "Preparing Data ...\n",
+      "Loading Trained Model ...\n",
       "Checking Model Outputs:\n",
       "\n",
       "Example 0 ========\n",
-      "\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
+      "\n",
       "Source Text (Input)        : <s> Baby sieht sich die Blätter am Zweig eines Baumes an . </s>\n",
-      "Target Text (Ground Truth) : <s> Baby looking at the leaves on a branch of a tree . </s>\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
+      "Target Text (Ground Truth) : <s> Baby looking at the leaves on a branch of a tree . </s>\n",
       "Model Output               : <s> A baby is looking at the leaves at a tree . </s>\n"
      ]
     },
@@ -4035,7 +3941,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.10"
+   "version": "3.10.9"
   }
  },
  "nbformat": 4,