code: def _get_module_flops_and_activation_sizes(
      module: torch.nn.Module,
      # pyre-ignore
      module_args: Optional[Tuple[Any, ...]] = None,
      # pyre-ignore
      module_kwargs: Optional[MutableMapping[str, Any]] = None,
  ) -> _ModuleSummaryData:
      # a mapping from module name to activation size tuple (in_size, out_size)
      activation_sizes: Dict[
          str,
          Tuple[Union[TUnknown, List[int]], Union[TUnknown, List[int]]],
      ] = {}
      # place activation size hooks on all modules + submodules
      activation_size_handles = _register_hooks(
          module,
          [
              (
                  _activation_size_hook(activation_sizes),
                  _HookType.FORWARD_HOOK,
              )
          ],
      )
  
      forward_timer_mapping: Dict[str, float] = {}
      forward_elapsed_times_sec: Dict[str, float] = {}
      forward_elapsed_time_handles = _register_hooks(
          module,
          [
              (
                  _forward_time_pre_hook(forward_timer_mapping),
                  _HookType.FORWARD_PRE_HOOK,
              ),
              (
                  _forward_time_hook(
                      forward_timer_mapping, forward_elapsed_times_sec
                  ),
                  _HookType.FORWARD_HOOK,
              ),
          ],
      )
  
      module.zero_grad()
  
      module_args = module_args or ()
      module_kwargs = module_kwargs or {}
      flops_forward = None
      flops_backward = None
      if not is_torch_version_geq_1_13():
          warnings.warn(
              "Please install PyTorch 1.13 or higher to compute FLOPs: https://pytorch.org/get-started/locally/"
          )
          module(*module_args, **module_kwargs)
          # detach activation size hook handles
          for hook_handle in activation_size_handles:
              hook_handle.remove()
      else:
          from torchtnt.utils.flops import FlopTensorDispatchMode
  
          with FlopTensorDispatchMode(module) as ftdm:
              # Count for forward flops (+ compute activation sizes)
              res = module(*module_args, **module_kwargs)
  
              # detach activation size hook handles
              for hook_handle in activation_size_handles:
                  hook_handle.remove()
  
              flops_forward = copy.deepcopy(ftdm.flop_counts)
              if isinstance(res, torch.Tensor):
                  # Count for backward flops
                  ftdm.reset()
                  res.mean().backward()
                  flops_backward = copy.deepcopy(ftdm.flop_counts)
              else:
                  warnings.warn(
                      "Backward FLOPs are only computed if module foward returns a tensor."
                  )
  
      # remove forward time elapsed handles
      for hook_handle in forward_elapsed_time_handles:
          hook_handle.remove()
  
      # convert module timings to ms
      forward_time_elapsed_ms = {}
      for module_name in forward_elapsed_times_sec:
          forward_time_elapsed_ms[module_name] = (
              forward_elapsed_times_sec[module_name] / 1000.0
          )
  
      # Reverting all the changes:
      module.zero_grad()
      module_summary_data = _ModuleSummaryData(
          flops_forward,
          flops_backward,
          activation_sizes,
          forward_time_elapsed_ms,
      )
      # TODO: Reverting BN: We also need to save status of BN running mean/var before running and revert those.
      return module_summary_data
  
imports: import copy
  import math
  import warnings
  from collections import defaultdict, deque
  from dataclasses import dataclass, field
  from enum import Enum
  from time import perf_counter
  from typing import Any, Callable, DefaultDict, Deque, Dict, List, MutableMapping, Optional, Tuple, Union
  import torch
  from torch.nn.parameter import UninitializedParameter
  from torch.utils._pytree import PyTree, tree_flatten
  from torch.utils.hooks import RemovableHandle
  from torchtnt.utils.version import is_torch_version_geq_1_13
  from typing_extensions import Literal
label_desc: Use image node_id 2 for calling a global function with example usage: _get_module_flops_and_activation_sizes(module, module_args, module_kwargs) and returns: module_summary_data