Update ttt_mlp_LayerNorm.py

ttt_mlp_LayerNorm.py

@@ -77,8 +77,8 @@ class TTT(nn.Module):
                 for param, grad in zip(self.mapping.parameters(), grads):
               
                     param -= 0.01 * grad
-            probe = self.probe(self.mapping(in_seq[:,seq,:]).detach())
-            outs.append(probe)
+            readout = self.mapping(self.probe(in_seq[:,seq,:])).detach()
+            outs.append(readout)
         out = torch.stack(outs, dim=1)
         
         return out