Michael Cusack
27 lines
612 B
Python
27 lines
612 B
Python
#!/usr/bin/env python3
|
|
#!/usr/bin/env python
|
|
"""Saves the model as a TorchScript.
|
|
|
|
The resulting file can be loaded in C++ code and then used for training or infrence with:
|
|
#include <torch/script.h>
|
|
torch::jit::Module module = torch::jit::load("model.pt")
|
|
"""
|
|
|
|
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()
|