Save model to file and generate images

1 files changed, 6 insertions, 2 deletions

diff --git a/mnist.py b/mnist.py
index 4466052..8ab8614 100644
--- a/mnist.py
+++ b/mnist.py
@@ -113,14 +113,18 @@ for epoch in range(num_epochs):
             print("Iteration: {}, Loss: {}, Accuracy: {}%".format(count, loss.data, accuracy))
 
 
+torch.save(model.state_dict(), "model.pth")
+print("Saved PyTorch Model State to model.pth")
+
+
 plt.plot(iteration_list, loss_list)
 plt.xlabel("No. of Iteration")
 plt.ylabel("Loss")
 plt.title("Iterations vs Loss")
-plt.show()
+plt.savefig("loss.png")
 
 plt.plot(iteration_list, accuracy_list)
 plt.xlabel("No. of Iteration")
 plt.ylabel("Accuracy")
 plt.title("Iterations vs Accuracy")
-plt.show()
+plt.savefig("accuracy.png")