Michael Cusack
22 lines
429 B
Python
22 lines
429 B
Python
#!/usr/bin/env python3
|
|
#!/usr/bin/env python
|
|
"""Saves the model as a TorchScript."""
|
|
|
|
import glob
|
|
import os
|
|
import sys
|
|
from typing import List
|
|
|
|
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_args)
|
|
torch.jit.save(torch.jit.script(model), "model.pt")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|