docstr: Test ModuleSummary callback in train
code: def test_module_summary_retrieve_module_summaries(self) -> None:
      """
      Test ModuleSummary callback in train
      """
  
      class Net(torch.nn.Module):
          def __init__(self):
              super(Net, self).__init__()
              self.l1 = torch.nn.Linear(2, 2)
              self.b1 = torch.nn.BatchNorm1d(2)
              self.l2 = torch.nn.Linear(2, 2)
  
          def forward(self, x):
              x = self.l1(x)
              x = self.b1(x)
              x = self.l2(x)
              return x
  
      my_module = Net()
  
      auto_unit = DummyAutoUnit(
          module=my_module,
      )
  
      module_summary_callback = ModuleSummary()
      summaries = module_summary_callback._retrieve_module_summaries(
          auto_unit
      )
  
      self.assertEqual(len(summaries), 1)
      ms = summaries[0]
      self.assertEqual(ms.module_name, "module")
      self.assertEqual(ms.module_type, "Net")
      self.assertTrue("l1" in ms.submodule_summaries)
      self.assertTrue("b1" in ms.submodule_summaries)
      self.assertTrue("l2" in ms.submodule_summaries)
  
file_path: tnt/tests/framework/callbacks/test_module_summary.py
imports: import unittest
  from unittest.mock import MagicMock
  import torch
  from torchtnt.framework._test_utils import DummyAutoUnit, DummyTrainUnit, generate_random_dataloader
  from torchtnt.framework.callbacks.module_summary import ModuleSummary
  from torchtnt.framework.state import EntryPoint, PhaseState, State
  from torchtnt.utils.version import is_torch_version_geq_1_13
label_desc: Use image node_id 2 for calling the ModuleSummaryTest obj's underlying member method code with example usage: obj.test_module_summary_retrieve_module_summaries() and returns: x