add find_unused_parameters param (#668)

* add find_unused_parameters param * fix default val
2021-04-06 21:29:38 +08:00 · 2021-04-06 21:29:38 +08:00 · a6e2114e32
parent 794caef23b
commit a6e2114e32
1 changed files with 3 additions and 1 deletions
--- a/tools/train.py
+++ b/tools/train.py
@ -70,7 +70,9 @@ def main(args):
        config, parameter_list=net.parameters())

    if config["use_data_parallel"]:
-        net = paddle.DataParallel(net)
+        find_unused_parameters = config.get("find_unused_parameters", False)
+        net = paddle.DataParallel(
+            net, find_unused_parameters=find_unused_parameters)

    # load model from checkpoint or pretrained model
    init_model(config, net, optimizer)