From c0954ee9ec700e0a0ebf4036a2bfccbcfaa7bce6 Mon Sep 17 00:00:00 2001 From: Marcin Chrzanowski Date: Sun, 2 May 2021 20:35:37 +0200 Subject: Save trained model --- src/experiment.py | 3 +++ 1 file changed, 3 insertions(+) (limited to 'src') 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) -- cgit v1.2.3