fix indentation issue and issues with save
@@ -13,6 +13,7 @@ except Exception as e:
raise ModuleNotFoundError('Need to install Pytorch: go to pytorch.org')
from collections import deque, namedtuple
+import pickle
import random
Batch = namedtuple("Batch", [
@@ -77,8 +78,6 @@ class TorchAgent(Agent):
self.history_tokens = opt['history_tokens']
self.history_dialog = opt['history_dialog']
self.history_replies = opt['history_replies']
- self.model = None
- self.optimizer = None
def share(self):
shared = super().share()
@@ -283,15 +282,18 @@ class TorchAgent(Agent):
"""
path = self.opt.get('model_file', None) if path is None else path
- if path and hasattr(self, 'model'):
+ if path and hasattr(self, 'model') and hasattr(self, 'optimizer'):
model = {}
model['model'] = self.model.state_dict()
model['optimizer'] = self.optimizer.state_dict()
- model['opt'] = self.opt
with open(path, 'wb') as write:
torch.save(model, write)
+ # save opt file
+ with open(path + ".opt", 'wb') as handle:
+ pickle.dump(self.opt, handle, protocol=pickle.HIGHEST_PROTOCOL)
+
def load(self, path):
"""Return opt and model states.
