[vits2] support ddp for net_dur_disc

wenet-e2e · Feb 21, 2024 · 64ee6fa · 64ee6fa
1 parent fb37f10
commit 64ee6fa
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/wetts/vits/train.py b/wetts/vits/train.py
@@ -203,6 +203,8 @@ def main():
     # latest work for others
     net_g = DDP(net_g, device_ids=[rank], find_unused_parameters=True)
     net_d = DDP(net_d, device_ids=[rank], find_unused_parameters=True)
+    if net_dur_disc:
+        net_dur_disc = DDP(net_dur_disc, device_ids=[rank], find_unused_parameters=True)
 
     try:
         _, _, _, epoch_str = task.load_checkpoint(