Fix minor issues with #21736 (#22074)

apaszke · facebook-github-bot · commit f1775796dde6 · 2019-06-24T15:18:26.000-07:00
Summary: cc mrshenli Pull Request resolved: #22074 Differential Revision: D15965376 Pulled By: mrshenli fbshipit-source-id: 50ff96de6390817d8ea52c04322c6bee3d649b32
diff --git a/torch/nn/parallel/distributed.py b/torch/nn/parallel/distributed.py
@@ -390,6 +390,8 @@ def __setstate__(self, state):
         # If serializable, then the process group should be the default one
         self.process_group = _get_default_group()
         super(DistributedDataParallel, self).__setstate__(state)
+        self.__dict__.setdefault('require_forward_param_sync', True)
+        self.__dict__.setdefault('require_backward_grad_sync', True)
         self._ddp_init_helper()
 
     def _check_default_group(self):
@@ -425,8 +427,10 @@ def no_sync(self):
         """
         old_require_backward_grad_sync = self.require_backward_grad_sync
         self.require_backward_grad_sync = False
-        yield
-        self.require_backward_grad_sync = old_require_backward_grad_sync
+        try:
+            yield
+        finally:
+            self.require_backward_grad_sync = old_require_backward_grad_sync
 
     def forward(self, *inputs, **kwargs):
         if self.require_forward_param_sync:
