diff --git a/Jupyter_Notebooks/Chapter_05_Natural_Language_Processing/02-Attention.ipynb b/Jupyter_Notebooks/Chapter_05_Natural_Language_Processing/02-Attention.ipynb
index b231a46f..305bce42 100644
--- a/Jupyter_Notebooks/Chapter_05_Natural_Language_Processing/02-Attention.ipynb
+++ b/Jupyter_Notebooks/Chapter_05_Natural_Language_Processing/02-Attention.ipynb
@@ -895,7 +895,8 @@
     "    # PEpos + k can be represented as a linear function of PEpos.\n",
     "    pe[:, 0::2] = torch.sin(position * div_term)\n",
     "    pe[:, 1::2] = torch.cos(position * div_term)\n",
-    "    pe = pe.unsqueeze(0).transpose(0, 1)\n",
+    "    # pe = pe.unsqueeze(0).transpose(0, 1)\n",
+    "    pe = pe.unsqueeze(0)\n",
     "    self.register_buffer('pe', pe)\n",
     "\n",
     "  def forward(self, x):\n",
@@ -910,7 +911,7 @@
     "      x: Tensor\n",
     "        Output is of the same shape as input with dropout and positional encodings\n",
     "    \"\"\"\n",
-    "    x = x + self.pe[:x.size(0), :]\n",
+    "    x = x + self.pe[0,:x.size(1), :]\n",
     "    return self.dropout(x)"
    ]
   },