Distributed fix for non-all reduce call.

866454e0 · Christian Sarofeen · 2fa4dbaf · 866454e0
Commit 866454e0 authored Apr 25, 2018 by Christian Sarofeen
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 1 deletion

apex/parallel/distributed.py apex/parallel/distributed.py +3 -1

No files found.
--- a/apex/parallel/distributed.py
+++ b/apex/parallel/distributed.py
@@ -26,7 +26,9 @@ def flat_dist_call(tensors, call, extra_args=None):
            call(coalesced, *extra_args)
        else:
            call(coalesced)
-        coalesced /= dist.get_world_size()
+        if call is dist.all_reduce:
+            coalesced /= dist.get_world_size()
+            
        for buf, synced in zip(bucket, _unflatten_dense_tensors(coalesced, bucket)):
            buf.copy_(synced)