Update litemixer.py

litemixer.py

@@ -24,14 +24,14 @@ class GatingUnit(nn.Module):
         self.proj_1 =  nn.Linear(dim,dim,bias=False)
         self.proj_2 =  nn.Linear(dim,dim,bias=False)
             
-        self.relu = nn.SiLU()
+        self.silu = nn.SiLU()
        
              	   
     def forward(self, x):
 
         u, v = x, x 
         u = self.proj_1(u)
-        u = self.relu(u)
+        u = self.silu(u)
         v = self.proj_2(v)
         g = u * v
         
@@ -45,14 +45,14 @@ class LiteMixerBlock(nn.Module):
         self.norm =  VectorDynamicTanh(dim)
        
 
-        self.context_interaction = nn.Sequential(
+        self.context_process = nn.Sequential(
            
             Rearrange('b n d -> b d n'),
             GatingUnit(num_patch),
             Rearrange('b d n -> b n d')
         )
 
-        self.token_interaction = GatingUnit(dim)
+        self.token_process = GatingUnit(dim)
         
 
     def forward(self, x):
@@ -61,7 +61,7 @@ class LiteMixerBlock(nn.Module):
     
         x = self.norm(x)
     
-        x = self.context_interaction(x)
+        x = self.context_process(x)
     
         x = x + residual
     
@@ -69,7 +69,7 @@ class LiteMixerBlock(nn.Module):
     
         x = self.norm(x)    
     
-        x = self.token_interaction(x)
+        x = self.token_process(x)
     
         x = x + residual