Fix Capacitron training (#2086)

coqui-ai · Nov 1, 2022 · 5307a22 · 5307a22
1 parent 5ccef6e
commit 5307a22
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/TTS/tts/models/base_tts.py b/TTS/tts/models/base_tts.py
@@ -344,7 +344,7 @@ def get_data_loader(
             loader = DataLoader(
                 dataset,
                 batch_size=config.eval_batch_size if is_eval else config.batch_size,
-                shuffle=False,  # shuffle is done in the dataset.
+                shuffle=True,  # if there is no other sampler
                 collate_fn=dataset.collate_fn,
                 drop_last=False,  # setting this False might cause issues in AMP training.
                 sampler=sampler,

diff --git a/TTS/utils/capacitron_optimizer.py b/TTS/utils/capacitron_optimizer.py
@@ -38,9 +38,9 @@ def step(self):
         self.param_groups = self.primary_optimizer.param_groups
         self.primary_optimizer.step()
 
-    def zero_grad(self):
-        self.primary_optimizer.zero_grad()
-        self.secondary_optimizer.zero_grad()
+    def zero_grad(self, set_to_none=False):
+        self.primary_optimizer.zero_grad(set_to_none)
+        self.secondary_optimizer.zero_grad(set_to_none)
 
     def load_state_dict(self, state_dict):
         self.primary_optimizer.load_state_dict(state_dict[0])