docstr: The function to call at the end of the constructor of each module.
  FIXME(fjr) The module may be passed to this function multiple times?
docstr_tok: ["The","function","to","call","at","the","end","of","the","constructor","of","each","module",".","FIXME","(","fjr",")","The","module","may","be","passed","to","this","function","multiple","times","?"]
code: def _post_init_method(self, module: torch.nn.Module, *args, **kwargs):
      """
      The function to call at the end of the constructor of each module.
      FIXME(fjr) The module may be passed to this function multiple times?
      """
      name_list = []
      for name, param in _named_params_with_replica(module):
          if type(param) is ColoParameter:
              continue
  
          split = name.rfind(".")
          if split >= 0:  # param in submodule
              module_name = name[:split]
              param_name = name[split + 1 :]
          else:
              module_name = ""  # param in current module
              param_name = name
          name_list.append((module_name, param_name))
  
      replaced_tensors = (
          dict()
      )  # record mapping between (torch.Tensor, ColoTensor) to distinguish the same reference
      for module_name, param_name in name_list:
          submodule = module.get_submodule(module_name)
          param = submodule.get_parameter(param_name)
          if param in replaced_tensors:
              colo_param = replaced_tensors[param]
          else:
              colo_param = _convert_to_coloparam(
                  param,
                  self._device,
                  self._dtype,
                  self._default_pg,
                  self._default_dist_spec,
              )
              replaced_tensors[param] = colo_param
          delattr(submodule, param_name)
          setattr(submodule, param_name, colo_param)
          colo_param.shared_param_modules.append(submodule)
  
      param_number = 0
      meta_param_number = 0
      buffer_number = 0
      meta_buffer_number = 0
  
      for param in module.parameters():
          param_number += 1
          meta_param_number += param.device.type == "meta"
  
      for buffer in module.buffers():
          buffer_number += 1
          meta_buffer_number += buffer.device.type == "meta"
  
      if meta_param_number > 0 and meta_param_number != param_number:
          raise ValueError(
              "Meta parameters and valued parameters can not  be in the same model"
          )
      if meta_buffer_number > 0 and meta_buffer_number != buffer_number:
          raise ValueError(
              "Meta buffers and valued buffers can not be in the same model"
          )
  
      if meta_buffer_number == 0:
          for buffer in module.buffers():
              buffer.data = buffer.data.to(device=self._device)
  
file_path: ColossalAI/colossalai/legacy/zero/gemini/colo_init_context.py
imports: from typing import Any, Iterator, Optional, Tuple, Union
  import torch
  from torch import nn
  from colossalai.legacy.tensor import ProcessGroup
  from colossalai.tensor import ColoParameter, ColoTensor
  from colossalai.utils.model.utils import InsertPostInitMethodToModuleSubClasses
label_desc: Use image node_id 4 for calling the ColoInitContext obj's underlying member method code with example usage: obj._post_init_method(module) without return types