Use 410m

app.py

@@ -13,14 +13,14 @@ def generate(prompt, seq_len=128, temperature=0.8, filter_thres=0.9):
     heads = 8
     flash_attn = True
 
-    model = PaLM(
-        num_tokens=num_tokens, dim=dim, depth=depth, dim_head=dim_head, heads=heads, flash_attn=flash_attn
-    ).to(device).eval()
-
     # model = PaLM(
-    #     num_tokens=50304, dim=1024, depth=24, dim_head=128, heads=8, flash_attn=False, qk_rmsnorm = False,
+    #     num_tokens=num_tokens, dim=dim, depth=depth, dim_head=dim_head, heads=heads, flash_attn=flash_attn
     # ).to(device).eval()
 
+    model = PaLM(
+        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)
     model.load_state_dict(checkpoint)