working

app.py

@@ -31,8 +31,8 @@ def main():
             depth=24,
         ),
     )        
-    model.net.token_emb.weight.requires_grad_(False)
-    model.net.to_logits.weight.requires_grad_(False)
+    # model.net.token_emb.weight.requires_grad_(False)
+    # model.net.to_logits.weight.requires_grad_(False)
     model.net.to_logits = nn.Sequential(
         nn.LayerNorm(f_emb),
         model.net.to_logits,
@@ -43,7 +43,7 @@ def main():
         wandb.watch(model)
     
     model.load_state_dict(torch.load('model.pt'))
-    optim = AdamW(model.parameters(), 2e-5)
+    optim = AdamW(model.parameters(), 5e-6)
 
     bs = 1
     kk = 2048
@@ -77,7 +77,7 @@ def main():
         if i % 1000 == 0:
             accelerator.wait_for_everyone()
             unwrapped_model = accelerator.unwrap_model(model)
-            b, n = 4, 512
+            b, n = 1, 2048
             init = torch.tensor([[50256]]*b).to(accelerator.device)
             prd = unwrapped_model.generate(init, n)
             prd = [dsx.decode(p) for p in prd]