fix bug with s2s save/load
@@ -145,6 +145,7 @@ class Seq2seqAgent(Agent):
# all instances may need some params
self.truncate = opt['truncate'] if opt['truncate'] > 0 else None
self.history = {}
+ self.states = {}
# check for cuda
self.use_cuda = not opt.get('no_cuda') and torch.cuda.is_available()
@@ -172,11 +173,10 @@ class Seq2seqAgent(Agent):
print('[ Using CUDA ]')
torch.cuda.set_device(opt['gpu'])
- self.states = {}
if opt.get('model_file') and os.path.isfile(opt['model_file']):
# load model parameters if available
print('Loading existing model params from ' + opt['model_file'])
- new_opt, states = self.load(opt['model_file'])
+ new_opt, self.states = self.load(opt['model_file'])
# override model-specific options with stored ones
opt = self.override_opt(new_opt)
@@ -576,6 +576,6 @@ class Seq2seqAgent(Agent):
def load(self, path):
"""Return opt and model states."""
with open(path, 'rb') as read:
- model = torch.load(read)
+ states = torch.load(read)
- return model['opt'], model
+ return states['opt'], states
