planner / templates/ml-project/train.py.tmpl

raw file 392 bytes 
"""Training entrypoint for {{name}}."""

import argparse

def main():
    parser = argparse.ArgumentParser(description="{{description}}")
    parser.add_argument("--epochs", type=int, default=10)
    parser.add_argument("--lr", type=float, default=1e-3)
    args = parser.parse_args()
    print(args)
    # TODO: load data, train model, save checkpoint

if __name__ == "__main__":
    main()