AllReduce bucket fix. (#186)

90017d3a · Samyam Rajbhandari · GitHub · 4cbfcc75 · 90017d3a
Unverified Commit 90017d3a authored Apr 11, 2020 by Samyam Rajbhandari Committed by GitHub Apr 11, 2020
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 0 deletions

deepspeed/pt/deepspeed_light.py deepspeed/pt/deepspeed_light.py +1 -0

No files found.
--- a/deepspeed/pt/deepspeed_light.py
+++ b/deepspeed/pt/deepspeed_light.py
@@ -840,6 +840,7 @@ class DeepSpeedLight(Module):
            if numel > numel_per_bucket:
                self.allreduce_and_copy(small_bucket)
                small_bucket = []
+                numel = 0
        if len(small_bucket) > 0:
            self.allreduce_and_copy(small_bucket)