class_docstr: This is a temporary extension of `LazyModuleMixin` to support passing keyword
  arguments to lazy module's `forward` method.
  
  The long-term plan is to upstream this feature to `LazyModuleMixin`. Please see
  https://github.com/pytorch/pytorch/issues/59923 for details.
  
  Please see `TestLazyModuleExtensionMixin`, which contains unit tests that ensure:
    * `LazyModuleExtensionMixin._infer_parameters` has source code parity with
      torch.nn.modules.lazy.LazyModuleMixin._infer_parameters, except that the former
      can accept keyword arguments.
    * `LazyModuleExtensionMixin._call_impl` has source code parity with
      `torch.nn.Module._call_impl`, except that the former can pass keyword arguments
      to forward pre hooks."
class_docstr_tok: ["This","is","a","temporary","extension","of","`","LazyModuleMixin","`","to","support","passing","keyword","arguments","to","lazy","module","'s","`","forward","`","method",".","The","long-term","plan","is","to","upstream","this","feature","to","`","LazyModuleMixin","`",".","Please","see","https",":","\/\/github.com\/pytorch\/pytorch\/issues\/59923","for","details",".","Please","see","`","TestLazyModuleExtensionMixin","`",",","which","contains","unit","tests","that","ensure",":","*","`","LazyModuleExtensionMixin._infer_parameters","`","has","source","code","parity","with","torch.nn.modules.lazy.LazyModuleMixin._infer_parameters",",","except","that","the","former","can","accept","keyword","arguments",".","*","`","LazyModuleExtensionMixin._call_impl","`","has","source","code","parity","with","`","torch.nn.Module._call_impl","`",",","except","that","the","former","can","pass","keyword","arguments","to","forward","pre","hooks",".","''"]
docstr: Infers the size and initializes the parameters according to the provided input batch.
  
  Given a module that contains parameters that were declared inferrable
  using :class:`torch.nn.parameter.ParameterMode.Infer`, runs a forward pass
  in the complete module using the provided input to initialize all the parameters
  as needed.
  The module is set into evaluation mode before running the forward pass in order
  to avoid saving statistics or calculating gradients
docstr_tok: ["Infers","the","size","and","initializes","the","parameters","according","to","the","provided","input","batch",".","Given","a","module","that","contains","parameters","that","were","declared","inferrable","using",":","class",":","`","torch.nn.parameter.ParameterMode.Infer","`",",","runs","a","forward","pass","in","the","complete","module","using","the","provided","input","to","initialize","all","the","parameters","as","needed",".","The","module","is","set","into","evaluation","mode","before","running","the","forward","pass","in","order","to","avoid","saving","statistics","or","calculating","gradients"]
code: def _infer_parameters(
      self: _LazyExtensionProtocol, module, args, kwargs
  ) -> None:
      r"""Infers the size and initializes the parameters according to the provided input batch.
  
      Given a module that contains parameters that were declared inferrable
      using :class:`torch.nn.parameter.ParameterMode.Infer`, runs a forward pass
      in the complete module using the provided input to initialize all the parameters
      as needed.
      The module is set into evaluation mode before running the forward pass in order
      to avoid saving statistics or calculating gradients
      """
      kwargs = kwargs if kwargs else {}
      module.initialize_parameters(*args, **kwargs)
      if module.has_uninitialized_params():
          raise RuntimeError(
              f"module {self._get_name()} has not been fully initialized"
          )
      module._initialize_hook.remove()
      module._load_hook.remove()
      delattr(module, "_initialize_hook")
      delattr(module, "_load_hook")
      if module.cls_to_become is not None:
          module.__class__ = module.cls_to_become
  
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 the LazyModuleExtensionMixin obj's underlying member method code with example usage: obj._infer_parameters(module, args, kwargs) without return types