diff options
Diffstat (limited to 'models/model.py')
| -rw-r--r-- | models/model.py | 2 |
1 files changed, 2 insertions, 0 deletions
diff --git a/models/model.py b/models/model.py index 5f079b8..eae15e3 100644 --- a/models/model.py +++ b/models/model.py @@ -174,6 +174,7 @@ class Model: print('{0:5d} loss: {1:.3f}'.format(iter_, loss)) self.rgb_pn.load_state_dict(checkpoint['model_state_dict']) self.optimizer.load_state_dict(checkpoint['optim_state_dict']) + self.scheduler.load_state_dict(checkpoint['sched_state_dict']) # Training start start_time = datetime.now() @@ -246,6 +247,7 @@ class Model: 'iter': self.curr_iter, 'model_state_dict': self.rgb_pn.state_dict(), 'optim_state_dict': self.optimizer.state_dict(), + 'sched_state_dict': self.scheduler.state_dict(), 'loss': loss, }, self._checkpoint_name) |