Move other logic after forward to take advantage of GPU skew

89fa152b · mcarilli · GitHub · 9d731777 · 89fa152b
Unverified Commit 89fa152b authored Sep 28, 2018 by mcarilli Committed by GitHub Sep 28, 2018
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
@@ -324,6 +324,8 @@ class DistributedDataParallel(Module):

        
    def forward(self, *inputs, **kwargs):
+        result = self.module(*inputs, **kwargs)
+        
        if not self.delay_allreduce:
            param_list = [param for param in self.module.parameters() if param.requires_grad]

@@ -354,4 +356,4 @@ class DistributedDataParallel(Module):

        self.callback_queued = False
        
-        return self.module(*inputs, **kwargs)
+        return result