unused params fix

Aske-Rosted · May 21, 2024 · 6fcd7d2 · 6fcd7d2
1 parent f5167f3
commit 6fcd7d2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/graphnet/models/components/embedding.py b/src/graphnet/models/components/embedding.py
@@ -84,7 +84,6 @@ def __init__(
         super().__init__()
 
         self.sin_emb = SinusoidalPosEmb(dim=seq_length, scaled=scaled)
-        self.aux_emb = nn.Embedding(2, seq_length // 2)
         self.sin_emb2 = SinusoidalPosEmb(dim=seq_length // 2, scaled=scaled)
 
         if n_features < 4:
@@ -93,6 +92,7 @@ def __init__(
                 f"{n_features} features."
             )
         elif n_features >= 6:
+            self.aux_emb = nn.Embedding(2, seq_length // 2)
             hidden_dim = 6 * seq_length
         else:
             hidden_dim = int((n_features + 0.5) * seq_length)