set find_unused_parameters=True in DDP

2019-06-18 12:13:14 +02:00
parent a59abedfb5
commit d82e5deeb1
2 changed files with 8 additions and 4 deletions
--- a/examples/run_squad.py
+++ b/examples/run_squad.py
@@ -907,7 +907,10 @@ def main():
    #     except ImportError:
    #         raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training.")

-        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank)
+        model = torch.nn.parallel.DistributedDataParallel(model,
+                                                          device_ids=[args.local_rank],
+                                                          output_device=args.local_rank,
+                                                          find_unused_parameters=True)
    elif n_gpu > 1:
        model = torch.nn.DataParallel(model)