need to unwrap

lucidrains · Oct 13, 2024 · 8cea21c · 8cea21c
1 parent 9599691
commit 8cea21c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/e2_tts_pytorch/trainer.py b/e2_tts_pytorch/trainer.py
@@ -244,7 +244,7 @@ def train(self, train_dataset, epochs, batch_size, num_workers=12, save_step=100
 
                     velocity_consistency_model = None
                     if self.need_velocity_consistent_loss and self.ema_model.initted:
-                        velocity_consistency_model = self.ema_model.ema_model
+                        velocity_consistency_model = self.accelerator.unwrap_model(self.ema_model).ema_model
 
                     loss, cond, pred, pred_data = self.model(
                         mel_spec,
@@ -262,7 +262,7 @@ def train(self, train_dataset, epochs, batch_size, num_workers=12, save_step=100
                     self.scheduler.step()
                     self.optimizer.zero_grad()
 
-                self.ema_model.update()
+                self.accelerator.unwrap_model(self.ema_model).update()
 
                 if self.accelerator.is_local_main_process:
                     logger.info(f"step {global_step+1}: loss = {loss.item():.4f}")