docstr: Test ModuleSummary callback in train
code: def test_module_summary_retrieve_module_summaries_module_inputs(
      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_inputs = {
          "module": ((torch.rand(2, 2, device=auto_unit.device),), {})
      }
      module_summary_callback = ModuleSummary(
          module_inputs=module_inputs
      )
      summaries = module_summary_callback._retrieve_module_summaries(
          auto_unit
      )
  
      self.assertEqual(len(summaries), 1)
      ms = summaries[0]
      self.assertEqual(ms.flops_forward, 16)
      self.assertEqual(ms.flops_backward, 24)
      self.assertEqual(ms.in_size, [2, 2])
      self.assertTrue(ms.forward_elapsed_time_ms != "?")
  
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 3 for calling the ModuleSummaryTest obj's underlying member method code with example usage: obj.test_module_summary_retrieve_module_summaries_module_inputs() and returns: x