Update normalizer.py

normalizer.py

@@ -2,11 +2,11 @@ import torch
 from torch import nn
 
 class GatingUnit(nn.Module):
-    def __init__(self, dim):
+    def __init__(self, dim, hidden_dim):
         super().__init__()
-        self.proj_1 =  nn.Linear(dim, dim)
-        self.proj_2 =  nn.Linear(dim, dim)
-        self.proj_3 = nn.Linear(dim, dim)     
+        self.proj_1 =  nn.Linear(dim, hidden_dim)
+        self.proj_2 =  nn.Linear(dim, hidden_dim)
+        self.proj_3 = nn.Linear(hidden_dim, dim)     
         self.silu = nn.SiLU()
                      	   
     def forward(self, x):
@@ -23,12 +23,12 @@ class GatingUnit(nn.Module):
         return out
 
 class NormalizerBlock(nn.Module):
-    def __init__(self, d_model, num_tokens):
+    def __init__(self, d_model, d_ffn, num_tokens):
         super().__init__()
                 
         self.norm_global = nn.LayerNorm(d_model * num_tokens, elementwise_affine=False)
         self.norm_local = nn.LayerNorm(d_model, elementwise_affine=False)                    
-        self.gating = GatingUnit(d_model)
+        self.gating = GatingUnit(d_model, d_ffn)
         
     def forward(self, x):
                   
@@ -53,7 +53,7 @@ class NormalizerBlock(nn.Module):
         return out
 
 class Normalizer(nn.Module):
-    def __init__(self, d_model, num_tokens, num_layers):
+    def __init__(self, d_model, d_ffn, num_tokens, num_layers):
         super().__init__()
         
         self.model = nn.Sequential(