code: def _activation_size_hook(
      activation_sizes: Dict[
          str,
          Tuple[Union[TUnknown, List[int]], Union[TUnknown, List[int]]],
      ],
      # pyre-ignore: Invalid type parameters [24]
  ) -> Callable[[str], Callable]:
      # pyre-ignore: Missing parameter annotation [2]
      def intermediate_hook(
          module_name: str,
      ) -> Callable[[torch.nn.Module, Any, Any], None]:
          # pyre-ignore
          def hook(_: torch.nn.Module, inp: Any, out: Any) -> None:
              if len(inp) == 1:
                  inp = inp[0]
              in_size = _parse_batch_shape(inp)
              out_size = _parse_batch_shape(out)
              activation_sizes[module_name] = (in_size, out_size)
  
          return hook
  
      return intermediate_hook
  
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 12 for calling a global function with example usage: _activation_size_hook(activation_sizes) and returns: intermediate_hook, hook