Update utils.py

jaywalnut310 · Jun 23, 2020 · e67f071 · e67f071
1 parent 00a482d
commit e67f071
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/utils.py b/utils.py
@@ -17,9 +17,11 @@
 def load_checkpoint(checkpoint_path, model, optimizer=None):
   assert os.path.isfile(checkpoint_path)
   checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')
-  iteration = checkpoint_dict['iteration']
-  learning_rate = checkpoint_dict['learning_rate']
-  if optimizer is not None:
+  if 'iteration' in checkpoint_dict.keys():
+    iteration = checkpoint_dict['iteration']
+  if 'learning_rate' in checkpoint_dict.keys():
+    learning_rate = checkpoint_dict['learning_rate']
+  if optimizer is not None and 'optimizer' in checkpoint_dict.keys():
     optimizer.load_state_dict(checkpoint_dict['optimizer'])
   saved_state_dict = checkpoint_dict['model']
   if hasattr(model, 'module'):