Update checkpoint.py

yebiny · Apr 9, 2021 · dbec79e · dbec79e
1 parent aa3ab7f
commit dbec79e
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/kospeech/checkpoint/checkpoint.py b/kospeech/checkpoint/checkpoint.py
@@ -74,7 +74,7 @@ def save(self):
             'optimizer': self.optimizer,
             'trainset_list': self.trainset_list,
             'validset': self.validset,
-            'epoch': self.epoch
+            'epoch': self.epoch,
         }
         torch.save(trainer_states, os.path.join(os.getcwd(), self.TRAINER_STATE_NAME))
         torch.save(self.model, os.path.join(os.getcwd(), self.MODEL_NAME))
@@ -111,9 +111,11 @@ def load(self, path):
                 model.flatten_parameters()
 
         return Checkpoint(
-            model=model, optimizer=resume_checkpoint['optimizer'], epoch=resume_checkpoint['epoch'],
+            model=model, 
+            optimizer=resume_checkpoint['optimizer'], 
+            epoch=resume_checkpoint['epoch'],
             trainset_list=resume_checkpoint['trainset_list'],
-            validset=resume_checkpoint['validset']
+            validset=resume_checkpoint['validset'],
         )
 
     def get_latest_checkpoint(self):