diff --git a/_modules/graphnet/models/components/embedding.html b/_modules/graphnet/models/components/embedding.html
index 06128b82d..9fcf13054 100644
--- a/_modules/graphnet/models/components/embedding.html
+++ b/_modules/graphnet/models/components/embedding.html
@@ -443,7 +443,6 @@ <h1 id="modules-graphnet-models-components-embedding--page-root">Source code for
         <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
 
         <span class="bp">self</span><span class="o">.</span><span class="n">sin_emb</span> <span class="o">=</span> <span class="n">SinusoidalPosEmb</span><span class="p">(</span><span class="n">dim</span><span class="o">=</span><span class="n">seq_length</span><span class="p">,</span> <span class="n">scaled</span><span class="o">=</span><span class="n">scaled</span><span class="p">)</span>
-        <span class="bp">self</span><span class="o">.</span><span class="n">aux_emb</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">Embedding</span><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="n">seq_length</span> <span class="o">//</span> <span class="mi">2</span><span class="p">)</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">sin_emb2</span> <span class="o">=</span> <span class="n">SinusoidalPosEmb</span><span class="p">(</span><span class="n">dim</span><span class="o">=</span><span class="n">seq_length</span> <span class="o">//</span> <span class="mi">2</span><span class="p">,</span> <span class="n">scaled</span><span class="o">=</span><span class="n">scaled</span><span class="p">)</span>
 
         <span class="k">if</span> <span class="n">n_features</span> <span class="o">&lt;</span> <span class="mi">4</span><span class="p">:</span>
@@ -452,7 +451,7 @@ <h1 id="modules-graphnet-models-components-embedding--page-root">Source code for
                 <span class="sa">f</span><span class="s2">"</span><span class="si">{</span><span class="n">n_features</span><span class="si">}</span><span class="s2"> features."</span>
             <span class="p">)</span>
         <span class="k">elif</span> <span class="n">n_features</span> <span class="o">&gt;=</span> <span class="mi">6</span><span class="p">:</span>
-
+            <span class="bp">self</span><span class="o">.</span><span class="n">aux_emb</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">Embedding</span><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="n">seq_length</span> <span class="o">//</span> <span class="mi">2</span><span class="p">)</span>
             <span class="n">hidden_dim</span> <span class="o">=</span> <span class="mi">6</span> <span class="o">*</span> <span class="n">seq_length</span>
         <span class="k">else</span><span class="p">:</span>
             <span class="n">hidden_dim</span> <span class="o">=</span> <span class="nb">int</span><span class="p">((</span><span class="n">n_features</span> <span class="o">+</span> <span class="mf">0.5</span><span class="p">)</span> <span class="o">*</span> <span class="n">seq_length</span><span class="p">)</span>