update save logic
@@ -282,13 +282,16 @@ class TorchAgent(Agent):
"""
path = self.opt.get('model_file', None) if path is None else path
- if path and hasattr(self, 'model') and hasattr(self, 'optimizer'):
- model = {}
- model['model'] = self.model.state_dict()
- model['optimizer'] = self.optimizer.state_dict()
-
+ if path:
+ states = {}
+ if hasattr(self, 'model'): # save model params
+ states['model'] = self.model.state_dict()
+ if hasattr(self, 'optimizer'): # save optimizer params
+ states['optimizer'] = self.optimizer.state_dict()
+
+ if states: # anything found to save?
with open(path, 'wb') as write:
- torch.save(model, write)
+ torch.save(states, write)
# save opt file
with open(path + ".opt", 'wb') as handle:
