fix(train): step lr schedulers at end of epoch (#291)

voicepaw · Apr 11, 2023 · 3af223e · 3af223e
1 parent 94a062b
commit 3af223e
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/src/so_vits_svc_fork/train.py b/src/so_vits_svc_fork/train.py
@@ -418,6 +418,11 @@ def training_step(self, batch: dict[str, torch.Tensor], batch_idx: int) -> None:
         optim_d.step()
         self.untoggle_optimizer(optim_d)
 
+        # end of epoch
+        if self.trainer.is_last_batch:
+            self.scheduler_g.step()
+            self.scheduler_d.step()
+
     def validation_step(self, batch, batch_idx):
         with torch.no_grad():
             self.net_g.eval()