Update smallformer.py

smallformer.py

@@ -5,6 +5,24 @@ import math
 
 
 
+class MLP(nn.Module):
+
+    def __init__(self,dim):
+        super().__init__()
+        self.proj_1 =  nn.Linear(dim,dim,bias=False)
+        self.proj_2 =  nn.Linear(dim,dim,bias=False)        
+        self.gelu = nn.GELU()
+       
+             	   
+    def forward(self, x):
+       
+        x = self.proj_1(x)
+        x = self.gelu(x)          
+        x = self.proj_2(x)
+        
+                          
+        return x
+
      
         
 
@@ -13,14 +31,14 @@ class LocalMappingUnit(nn.Module):
         super().__init__()
         
            
-        self.mapping = nn.PReLU(dim)
+        self.mapping = MLP(dim)
         self.norm = nn.LayerNorm(dim,elementwise_affine=False)
       
              	   
     def forward(self, x):
     
         x = self.norm(x) 
-        x = self.mapping(x[-1])   	
+        x = self.mapping(x)   	
       
         return x