Save trained model

1 files changed, 3 insertions, 0 deletions

diff --git a/src/experiment.py b/src/experiment.py
index b508438..f8a219a 100644
--- a/src/experiment.py
+++ b/src/experiment.py
@@ -4,6 +4,7 @@ import time
 
 import pandas as pd
 import matplotlib.pyplot as plt
+import torch
 
 from runner import Runner
 
@@ -35,6 +36,8 @@ class Experiment:
         plt.legend()
         plt.savefig(self.dir_path('accuracies.png'))
 
+        torch.save(self.runner.net.state_dict(), self.dir_path('net.pt'))
+
     def dir_path(self, file):
         return '{}/{}'.format(self.dirname, file)