code: def lazy_apply(
      module: torch.nn.Module, fn: Callable[[torch.nn.Module], None]
  ) -> torch.nn.Module:
      """Attaches a function to a module, which will be applied recursively to every
      submodule (as returned by `.children()`) of the module as well as the module itself
      right after the first forward pass (i.e. after all submodules and parameters have
      been initialized).
  
      Typical use includes initializing the numerical value of the parameters of a lazy
      module (i.e. modules inherited from `LazyModuleMixin`).
  
      NOTE:
          `lazy_apply()` can be used on both lazy and non-lazy modules.
  
      Args:
          module (torch.nn.Module): module to recursively apply `fn` on.
          fn (Callable[[torch.nn.Module], None]): function to be attached to `module` and
              later be applied to each submodule of `module` and the `module` itself.
  
      Returns:
          torch.nn.Module: `module` with `fn` attached.
  
      Example::
  
          @torch.no_grad()
          def init_weights(m):
              print(m)
              if type(m) == torch.nn.LazyLinear:
                  m.weight.fill_(1.0)
                  print(m.weight)
  
          linear = torch.nn.LazyLinear(2)
          lazy_apply(linear, init_weights)  # doesn't run `init_weights` immediately
          input = torch.randn(2, 10)
          linear(input)  # runs `init_weights` only once, right after first forward pass
  
          seq = torch.nn.Sequential(torch.nn.LazyLinear(2), torch.nn.LazyLinear(2))
          lazy_apply(seq, init_weights)  # doesn't run `init_weights` immediately
          input = torch.randn(2, 10)
          seq(input)  # runs `init_weights` only once, right after first forward pass
      """
  
      if not hasattr(module, "_functions_to_lazy_apply"):
          module._functions_to_lazy_apply = []
      if not hasattr(module, "_lazy_apply_hook"):
          module._lazy_apply_hook = module.register_forward_hook(
              _apply_functions_after_first_forward
          )
      module._functions_to_lazy_apply.append(fn)
      return module
  
file_path: torchrec/torchrec/modules/lazy_extension.py
imports: import functools
  import inspect
  from typing import Any, Callable
  import torch
  import torch.utils.hooks
  from torch.nn.modules.lazy import _LazyProtocol, LazyModuleMixin
  from torch.nn.modules.module import _global_backward_hooks, _global_forward_hooks, _global_forward_pre_hooks
label_desc: Use image node_id 2 for calling a global function with example usage: lazy_apply(module, fn) and returns: module