update DDP call

Signed-off-by: Jason <[email protected]>
NVIDIA · Feb 13, 2020 · dabaea2 · dabaea2
1 parent f1a57bb
commit dabaea2
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/nemo/backends/pytorch/actions.py b/nemo/backends/pytorch/actions.py
@@ -1239,7 +1239,9 @@ def train(
 
                         # By default, disable broadcast_buffers. This disables batch norm synchronization on forward
                         # pass
-                        pmodule = DDP(pmodule, device_ids=[self.local_rank], broadcast_buffers=False)
+                        pmodule = DDP(
+                            pmodule, device_ids=[self.local_rank], broadcast_buffers=False, find_unused_parameters=True
+                        )
 
                     # # Convert batchnorm modules to synced if applicable
                     # if synced_batchnorm and isinstance(pmodule, torch.nn.Module):