fix typo mha

Author: vpj

docs/transformers/mha.html

@@ -208,7 +208,7 @@ <h2>Prepare for multi-head attention</h2>
                 <a href='#section-10'>#</a>
             </div>
             <p>Output has shape <code  class="highlight"><span></span><span class="p">[</span><span class="n">seq_len</span><span class="p">,</span> <span class="n">batch_size</span><span class="p">,</span> <span class="n">heads</span><span class="p">,</span> <span class="n">d_k</span><span class="p">]</span></code>
- or <code  class="highlight"><span></span><span class="p">[</span><span class="n">batch_size</span><span class="p">,</span> <span class="n">d_model</span><span class="p">]</span></code>
+ or <code  class="highlight"><span></span><span class="p">[</span><span class="n">batch_size</span><span class="p">,</span> <span class="n">heads</span><span class="p">,</span> <span class="n">d_model</span><span class="p">]</span></code>
  </p>
 
         </div>

labml_nn/transformers/mha.py

@@ -62,7 +62,7 @@ def forward(self, x: torch.Tensor):
         # Split last dimension into heads
         x = x.view(*head_shape, self.heads, self.d_k)
 
-        # Output has shape `[seq_len, batch_size, heads, d_k]` or `[batch_size, d_model]`
+        # Output has shape `[seq_len, batch_size, heads, d_k]` or `[batch_size, heads, d_model]`
         return x