Restore `last_epoch` of the scheduler

This commit is contained in:
Eren Gölge 2021-07-17 01:02:47 +02:00
parent aacbb3ed77
commit 9af42f7886
1 changed files with 5 additions and 2 deletions

View File

@ -268,9 +268,13 @@ class Trainer:
self.config, args.restore_path, self.model, self.optimizer, self.scaler
)
# setup scheduler
# setup scheduler
self.scheduler = self.get_scheduler(self.model, self.config, self.optimizer)
if self.args.continue_path:
self.scheduler.last_epoch = self.restore_step
# DISTRUBUTED
if self.num_gpus > 1:
self.model = DDP_th(self.model, device_ids=[args.rank], output_device=args.rank)
@ -291,7 +295,6 @@ class Trainer:
Returns:
nn.Module: initialized model.
"""
# TODO: better model setup
try:
model = setup_vocoder_model(config)
except ModuleNotFoundError: