fix ckpt resume bug

Author: MaureenZOU

trainer/default_trainer.py

@@ -227,7 +227,7 @@ def train(self):
 
         train_prev_logged_time = datetime.now()
         for epoch in range(self.train_params['start_epoch_idx'], num_epochs):
-            self.train_params['start_epoch_idx'] = epoch
+            self.train_params['current_epoch_idx'] = epoch
             logger.info(f"Start epoch: {epoch} training.")
 
             epoch_start_time = datetime.now()

trainer/utils_trainer.py

@@ -77,10 +77,15 @@ def save_checkpoint(self, tag):
         tag = str(tag).zfill(8)
         logger.warning('Saving checkpoint...')
 
-        self.train_params['current_batch_idx'] = self.train_params['current_batch_idx'] + 1
-        if self.train_params['current_batch_idx'] == self.train_params['updates_per_epoch']:
-            self.train_params['resume_batch_idx'] = 0
-            self.train_params['resume_epoch_idx'] += 1
+        resume_epoch_idx = self.train_params['current_epoch_idx']
+        resume_batch_idx = self.train_params['current_epoch_idx'] + 1
+
+        if resume_batch_idx == self.train_params['updates_per_epoch']:
+            self.train_params['start_batch_idx'] = 0
+            self.train_params['start_epoch_idx'] = resume_epoch_idx + 1
+        else:
+            self.train_params['start_batch_idx'] = resume_batch_idx
+            self.train_params['start_epoch_idx'] = resume_epoch_idx
 
         save_dir = os.path.join(self.save_folder, tag)