Update tensormapper.py

tensormapper.py

@@ -13,12 +13,12 @@ class VecDyT(nn.Module):
         return x
                 
 class GatingUnit(nn.Module):
-    def __init__(self,dim):
+    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.proj_1 = nn.Linear(dim, dim, bias = False)
+        self.proj_2 = nn.Linear(dim, dim, bias = False)
             
         self.gelu = nn.GELU()
                     	   
@@ -37,9 +37,9 @@ class TTT(nn.Module):
         
         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)
+        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:
        
@@ -52,14 +52,14 @@ class TTT(nn.Module):
             label_view = state
             loss = nn.functional.mse_loss(self.mapping(train_view), label_view)
             grads = torch.autograd.grad(
-                loss, self.mapping.parameters(),create_graph=True)
+                loss, self.mapping.parameters(), create_graph = True)
             with torch.no_grad():
                 for param, grad in zip(self.mapping.parameters(), grads):
               
                     param -= 0.01 * grad
             readout = self.mapping(self.Probe(in_seq[:,seq,:])).detach()
             outs.append(readout)
-        out = torch.stack(outs, dim=1)
+        out = torch.stack(outs, dim = 1)
         
         return out 
             
@@ -104,4 +104,4 @@ class TensorMapper(nn.Module):
 
     def forward(self, x):
        
-        return self.model(x)
+        return self.model(x)