From 0751bcce7e91ca3c5f8fb575c16435e8237c1ec9 Mon Sep 17 00:00:00 2001
From: Marcin Kostrzewski <komar86311@gmail.com>
Date: Sun, 24 Apr 2022 22:23:53 +0200
Subject: [PATCH] Model saving

---
 train_model.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/train_model.py b/train_model.py
index bdf091a..21bf880 100644
--- 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.")