Update modeling_tta.py

modeling_tta.py

@@ -49,7 +49,7 @@ def get_hadamard_matrix(d: int, dtype=torch.float32) -> torch.Tensor:
     return ortho.contiguous()
 
 class HadamardMLP(nn.Module):
-    def __init__(self, config: CustomConfig):
+    def __init__(self, config: TTAConfig):
         super().__init__()
         self.dim = config.hidden_size
         self.scale1 = nn.Parameter(torch.ones(self.dim))
@@ -68,7 +68,7 @@ class HadamardMLP(nn.Module):
         return out
 
 class SwiGLUMLP(nn.Module):
-    def __init__(self, config: CustomConfig):
+    def __init__(self, config: TTAConfig):
         super().__init__()
         self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
         self.up_proj   = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
@@ -78,7 +78,7 @@ class SwiGLUMLP(nn.Module):
         return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
 
 class EngramMemory(nn.Module):
-    def __init__(self, config: CustomConfig):
+    def __init__(self, config: TTAConfig):
         super().__init__()
         self.dim = config.hidden_size
         self.num_entries = config.engram_entries