Use 410m

app.py

@@ -21,7 +21,7 @@ def generate(prompt, seq_len=128, temperature=0.8, filter_thres=0.9):
         num_tokens=50304, dim=1024, depth=24, dim_head=128, heads=8, flash_attn=False, qk_rmsnorm = False,
     ).to(device).eval()
 
-    checkpoint = torch.load('./palm_1b_8k_v0.pt', map_location=device)
+    checkpoint = torch.load('./palm_410m_8k_v0.pt', map_location=device)
     model.load_state_dict(checkpoint)
 
     tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")