Save model to file and generate images

4 files changed, 6 insertions, 2 deletions

diff --git a/accuracy.png b/accuracy.png
new file mode 100644
index 0000000..4d04a64
--- /dev/null
+++ b/accuracy.png
diff --git a/loss.png b/loss.png
new file mode 100644
index 0000000..8a9eca4
--- /dev/null
+++ b/loss.png
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")
diff --git a/model.pth b/model.pth
index 195525e..d0a1381 100644
--- a/model.pth
+++ b/model.pth