Ensures that only GPT model is in training mode during training

coqui-ai · Nov 16, 2023 · b87a665 · b87a665
1 parent 7e4375d
commit b87a665
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/TTS/tts/layers/xtts/trainer/gpt_trainer.py b/TTS/tts/layers/xtts/trainer/gpt_trainer.py
@@ -318,9 +318,10 @@ def eval_step(self, batch, criterion):
         batch["cond_idxs"] = None
         return self.train_step(batch, criterion)
 
-    def on_epoch_start(self, trainer):  # pylint: disable=W0613
-        # guarante that dvae will be in eval mode after .train() on evaluation end
-        self.dvae = self.dvae.eval()
+    def on_train_epoch_start(self, trainer):
+        trainer.model.eval() # the whole model to eval
+        # put gpt model in training mode
+        trainer.model.xtts.gpt.train()
 
     def on_init_end(self, trainer):  # pylint: disable=W0613
         # ignore similarities.pth on clearml save/upload

diff --git a/requirements.txt b/requirements.txt
@@ -27,7 +27,7 @@ pandas>=1.4,<2.0
 # deps for training
 matplotlib==3.7.*
 # coqui stack
-trainer
+trainer>=0.0.32
 # config management
 coqpit>=0.0.16
 # chinese g2p deps