Model saving

2022-04-24 22:23:53 +02:00 · 2022-04-24 22:23:53 +02:00 · 0751bcce7e
commit 0751bcce7e
parent d95bf0d73d
1 changed files with 3 additions and 0 deletions
--- a/train_model.py
+++ b/train_model.py
@ -109,3 +109,6 @@ for t in range(epochs):
    train(train_dataloader, model, loss_fn, optimizer)
    test(test_dataloader, model, loss_fn)
 print("Done!")
+
+torch.save(model.state_dict(), './model_out')
+print("Model saved in ./model_out file.")