fn

Author: NihalHarish

smdebug/pytorch/hook.py

@@ -155,7 +155,7 @@ def forward_hook(self, module, inputs, outputs):
             return
 
         # _get_name() is implemented by nn.Module which is the base class for all modules
-        module_name = module._get_name()
+        module_name = module._module_name
         # This overwhelms the logs; turn back on if you really need it
         # logger.debug("Processing the global step {0} for module {1}".format(self.step, module_name))
 
@@ -198,6 +198,14 @@ def register_hook(self, module):
         # for compatibility with ZCC patches which call this
         self.register_module(module)
 
+    @staticmethod
+    def _add_module_name(module, module_name):
+        if isinstance(module, torch.nn.parallel.data_parallel.DataParallel):
+            module.module._module_name = module_name
+        else:
+            module._module_name = module_name
+        return module
+
     def register_module(self, module):
         """
         This function registers the forward hook. If user wants to register the hook
@@ -216,9 +224,9 @@ def register_module(self, module):
 
         for name, submodule in module.named_modules():
             assert submodule not in self.module_set, f"Don't register module={module} twice"
-            submodule._module_name = name
+            Hook._add_module_name(submodule, name)
             self.module_set.add(submodule)
-        module._module_name = module._get_name()
+        Hook._add_module_name(module, module._get_name())
         self.module_set.add(module)
 
         # Use `forward_pre_hook` for the entire net