Model args in save script

2023-08-04 17:07:41 +07:00
parent fd5e2cc7bc
commit f67185958b
1 changed files with 11 additions and 2 deletions
@@ -16,9 +16,18 @@ import torch
 from model import ModelArgs, Transformer
 def main() -> None:
-    model_args = ModelArgs(dim=512, n_layers=6, n_heads=8, vocab_size=32000)
+    model = Transformer(
-    model = Transformer(model_args)
+        ModelArgs(
            dim=288,
            n_layers=6,
            n_heads=6,
            multiple_of=32,
            dropout=0.0,
            vocab_size=32000,
        )
    )
    torch.jit.save(torch.jit.script(model), "model.pt")