From d0015a0ca43c6afc288cee985f8ae740de0c56bd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stardust=C2=B7=E5=87=8F?= Date: Sun, 3 Sep 2023 13:38:32 +0800 Subject: [PATCH] Update train_ms.py --- train_ms.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/train_ms.py b/train_ms.py index 4284a19..55f2bb5 100644 --- a/train_ms.py +++ b/train_ms.py @@ -147,10 +147,10 @@ def run(rank, n_gpus, hps): eps=hps.train.eps) else: optim_dur_disc = None - net_g = DDP(net_g, device_ids=[rank]) - net_d = DDP(net_d, device_ids=[rank]) + 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 is not None: - net_dur_disc = DDP(net_dur_disc, device_ids=[rank]) + net_dur_disc = DDP(net_dur_disc, device_ids=[rank], find_unused_parameters=True) pretrain_dir = None if pretrain_dir is None: