From f1a664e0d476d50c91701e8a2e283f1bfc23b872 Mon Sep 17 00:00:00 2001 From: Fangjun Kuang Date: Sat, 6 Aug 2022 12:52:25 +0800 Subject: [PATCH] Set find_unused_parameters optionally. --- egs/aishell/ASR/pruned_transducer_stateless3/train.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/egs/aishell/ASR/pruned_transducer_stateless3/train.py b/egs/aishell/ASR/pruned_transducer_stateless3/train.py index ce29abb74..feaef5cf6 100755 --- a/egs/aishell/ASR/pruned_transducer_stateless3/train.py +++ b/egs/aishell/ASR/pruned_transducer_stateless3/train.py @@ -1038,7 +1038,16 @@ def run(rank, world_size, args): model.to(device) if world_size > 1: logging.info("Using DDP") - model = DDP(model, device_ids=[rank], find_unused_parameters=True) + if params.datatang_prob > 0: + find_unused_parameters = True + else: + find_unused_parameters = False + + model = DDP( + model, + device_ids=[rank], + find_unused_parameters=find_unused_parameters, + ) optimizer = Eve(model.parameters(), lr=params.initial_lr)