v19: Fix tied weights and lm_head loading

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b73a2f8817a80324e50dcbdb5de0efd76813577c1d48700f60c6b6f3f1f2a3bc
-size 182717432
+oid sha256:5235b319424854b28c715719aad9870d4ffc4e97e986f7fe4803e1023a7e71d8
+size 182717408

modeling_trm_text_ism.py

@@ -39,7 +39,7 @@ class TRMBlock(nn.Module):
         self.res, self.n1, self.n2 = config.residual_scale, nn.RMSNorm(config.dim), nn.RMSNorm(config.dim)
         self.attn, self.mlp = TRMAttention(config), SwiGLUMLP(config)
         self.ag, self.mg = nn.Parameter(torch.ones(config.dim)), nn.Parameter(torch.ones(config.dim))
-    def forward(self, x, mask, c, s): 
+    def forward(self, x, mask, c, s):
         x = x + self.res * torch.sigmoid(self.ag) * self.attn(self.n1(x), mask, c, s)
         return x + self.res * torch.sigmoid(self.mg) * self.mlp(self.n2(x))
 
@@ -57,8 +57,8 @@ class TRMTextISMForCausalLM(PreTrainedModel):
         self.post_init()
     def get_input_embeddings(self): return self.token_emb
     def get_output_embeddings(self): return self.lm_head
-    def tie_weights(self): self.lm_head.weight = self.token_emb.weight
-    def forward(self, input_ids, attention_mask=None, **kwargs):
+    def tie_weights(self, *args, **kwargs): self.lm_head.weight = self.token_emb.weight
+    def forward(self, input_ids, **kwargs):
         B, S = input_ids.shape
         x, m = self.token_emb(input_ids), torch.tril(torch.ones(S, S, device=input_ids.device)).bool().unsqueeze(0).expand(B, -1, -1)
         for _ in range(self.config.recurrence_steps): x = self.block(x, m, self.rope_cos, self.rope_sin)