DDP bug fix

Signed-off-by: Jason <[email protected]>
NVIDIA · May 18, 2020 · fa6553f · fa6553f
1 parent d62f021
commit fa6553f
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/nemo/backends/pytorch/actions.py b/nemo/backends/pytorch/actions.py
@@ -1334,6 +1334,7 @@ def save_state_to(self, path):
             for module in AppState().modules:
                 key = module.unique_instance_id
                 num_trainable_weights = module.num_weights
+                self.ddp_module_dict[key] = module
                 if not isinstance(module, DDP) and isinstance(module, torch.nn.Module) and num_trainable_weights > 0:
                     # gpf = 1
                     # if gradient_predivide: