fix value pe double rotation

Author: vpj

docs/transformers/rope/value_pe/index.html

@@ -412,7 +412,7 @@ <h2>Multi-head attention with rotary positional embeddings</h2>
 
         </div>
         <div class='code'>
-            <div class="highlight"><pre><span class="lineno">234</span>        <span class="n">x</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">einsum</span><span class="p">(</span><span class="s2">&quot;ijbh,jbhd-&gt;ibhd&quot;</span><span class="p">,</span> <span class="n">attn</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">value_rotary_pe</span><span class="p">(</span><span class="n">value</span><span class="p">))</span></pre></div>
+            <div class="highlight"><pre><span class="lineno">234</span>        <span class="n">x</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">einsum</span><span class="p">(</span><span class="s2">&quot;ijbh,jbhd-&gt;ibhd&quot;</span><span class="p">,</span> <span class="n">attn</span><span class="p">,</span> <span class="n">value</span><span class="p">)</span></pre></div>
         </div>
     </div>
     <div class='section' id='section-21'>

labml_nn/transformers/rope/value_pe/__init__.py

@@ -231,7 +231,7 @@ def forward(self, *,
 
         # Multiply by values
         # $$\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_k}}\Bigg)V$$
-        x = torch.einsum("ijbh,jbhd->ibhd", attn, self.value_rotary_pe(value))
+        x = torch.einsum("ijbh,jbhd->ibhd", attn, value)
 
         # Rotate in the opposite direction so that each embedding hold the relative positions
         x = self.value_reverse_rotary_pe(x)