code: def init_weights(module):
      if type(module) in (th.nn.Linear, th.nn.Embedding):
          th.nn.init.trunc_normal_(
              module.weight, mean=0, std=0.02, a=-0.04, b=0.04
          )
          if type(module) == th.nn.Linear and module.bias is not None:
              module.bias.data.zero_()
      elif isinstance(module, th.nn.LayerNorm):
          module.bias.data.zero_()
          module.weight.data.fill_(1.0)
      else:
          logging.debug(
              f"Not performing custom initialization for {type(module)}"
          )
  
imports: import logging
  import torch
  from ..attention_cell import gen_self_attn_mask
  from .transformer import TransformerEncoderLayer
  from ..layers import get_activation, sequence_mask
  from ...utils.registry import Registry
  from ...models.bert import bert_cfg_reg
label_desc: Use image node_id 1 for calling a global function with example usage: init_weights(module) without return types