support multiple GPU training

coqui-ai · Dec 8, 2023 · b6e9296 · b6e9296
1 parent c99e885
commit b6e9296
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/TTS/tts/layers/xtts/trainer/gpt_trainer.py b/TTS/tts/layers/xtts/trainer/gpt_trainer.py
@@ -321,7 +321,10 @@ def eval_step(self, batch, criterion):
     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()
+        if hasattr(trainer.model, "module") and hasattr(trainer.model.module, "xtts"):
+            trainer.model.module.xtts.gpt.train()
+        else:
+            trainer.model.xtts.gpt.train()
 
     def on_init_end(self, trainer):  # pylint: disable=W0613
         # ignore similarities.pth on clearml save/upload
@@ -387,7 +390,8 @@ def get_data_loader(
             else:
                 loader = DataLoader(
                     dataset,
-                    batch_sampler=sampler,
+                    sampler=sampler,
+                    batch_size = config.eval_batch_size if is_eval else config.batch_size,
                     collate_fn=dataset.collate_fn,
                     num_workers=config.num_eval_loader_workers if is_eval else config.num_loader_workers,
                     pin_memory=False,