class_docstr: BERT-based encoder for queries or blocks used for learned information retrieval.
class_docstr_tok: ["BERT-based","encoder","for","queries","or","blocks","used","for","learned","information","retrieval","."]
code: def __init__(
      self, ict_head_size, num_tokentypes=2, parallel_output=True
  ):
      super(IREncoderBertModel, self).__init__()
      args = get_args()
  
      self.ict_head_size = ict_head_size
      self.parallel_output = parallel_output
      init_method = init_method_normal(args.init_method_std)
      scaled_init_method = scaled_init_method_normal(
          args.init_method_std, args.num_layers
      )
  
      (
          self.language_model,
          self._language_model_key,
      ) = get_language_model(
          num_tokentypes=num_tokentypes,
          add_pooler=True,
          encoder_attn_mask_type=AttnMaskType.padding,
          init_method=init_method,
          scaled_init_method=scaled_init_method,
      )
  
      self.ict_head = get_linear_layer(
          args.hidden_size, ict_head_size, init_method
      )
      self._ict_head_key = "ict_head"
  
imports: import os
  import torch
  from megatron import get_args, print_rank_0
  from megatron.checkpointing import get_checkpoint_tracker_filename, get_checkpoint_name
  from megatron.model import BertModel
  from .module import MegatronModule
  from megatron.core import mpu
  from megatron.model.enums import AttnMaskType
  from megatron.model.utils import get_linear_layer
  from megatron.model.utils import init_method_normal
  from megatron.model.language_model import get_language_model
  from megatron.model.utils import scaled_init_method_normal
  from megatron.model.bert_model import bert_extended_attention_mask, bert_position_ids
label_desc: Use image node_id 1 to create a new IREncoderBertModel object from inherited base classes: MegatronModule with example: obj = IREncoderBertModel(ict_head_size, num_tokentypes, parallel_output)