-except FileNotFoundError:
- print('Starting from scratch.')
-
-except:
- print('Error when loading the checkpoint.')
- exit(1)
+else:
+ try:
+ checkpoint = torch.load(args.checkpoint_name, map_location = device)
+ nb_epochs_finished = checkpoint['nb_epochs_finished']
+ model.load_state_dict(checkpoint['model_state'])
+ optimizer.load_state_dict(checkpoint['optimizer_state'])
+ log_string(f'Checkpoint loaded with {nb_epochs_finished} epochs finished.')
+
+ except FileNotFoundError:
+ log_string('Starting from scratch.')
+
+ except:
+ log_string('Error when loading the checkpoint.')
+ exit(1)