code: def get_module_output_dimension(
      module: Union[
          Callable[[torch.Tensor], torch.Tensor], torch.nn.Module
      ],
      in_features: int,
  ) -> int:
      input = torch.zeros(1, in_features)
      output = module(input)
      return output.size(-1)
  
file_path: torchrec/torchrec/modules/utils.py
imports: import copy
  from typing import Callable, Iterable, Tuple, Union
  import torch
label_desc: Use image node_id 2 for calling a global function with example usage: get_module_output_dimension(module, in_features) and returns: output