too many layernorms

lucidrains · Nov 9, 2021 · c8d1ede · c8d1ede
1 parent 01f75d7
commit c8d1ede
Showing 1 changed file with 1 addition and 4 deletions.
diff --git a/hourglass_transformer_pytorch/hourglass_transformer_pytorch.py b/hourglass_transformer_pytorch/hourglass_transformer_pytorch.py
@@ -194,10 +194,7 @@ def __init__(
             norm_out = True
         )
 
-        self.to_logits = nn.Sequential(
-            nn.LayerNorm(dim),
-            nn.Linear(dim, num_tokens)
-        )
+        self.to_logits = nn.Linear(dim, num_tokens)
 
     def forward(self, x):
         device = x.device