Update tensormapper.py

tensormapper.py

@@ -7,17 +7,17 @@ class VecDyT(nn.Module):
     def __init__(self, input_shape):
     
         super().__init__()
-        
-           
+                   
         self.alpha = nn.Parameter(torch.randn(input_shape))
        
-
     def forward(self, x):
         x = torch.tanh(self.alpha * x)
         return x
         
+        
 class GatingUnit(nn.Module):
     def __init__(self,dim):
+        
         super().__init__()
 
         self.proj_1 =  nn.Linear(dim,dim,bias=False)
@@ -36,21 +36,16 @@ class GatingUnit(nn.Module):
         
         return g
 
-class TTT(nn.Module):
-   
-
+class TTT(nn.Module):   
     def __init__(self, dim: int):
+        
         super(TTT, self).__init__()
-       
-     
-
-       
+            
        
         self.mapping = nn.Linear(dim,dim,bias=False)
         self.State =  nn.Linear(dim,dim,bias=False)
         self.Probe =  nn.Linear(dim,dim,bias=False)
-       
-        
+               
        
     def forward(self, in_seq: Tensor) -> Tensor:
 
@@ -78,8 +73,8 @@ class TTT(nn.Module):
 
 
 class TensorMapperBlock(nn.Module):
-
     def __init__(self, dim, num_patch):
+        
         super().__init__()
         
         self.norm_1 =  VecDyT(dim)