class_docstr: Language model embeddings.
  
  Arguments:
      hidden_size: hidden size
      vocab_size: vocabulary size
      max_sequence_length: maximum size of sequence. This
                           is used for positional embedding
      embedding_dropout_prob: dropout probability for embeddings
      init_method: weight initialization method
      num_tokentypes: size of the token-type embeddings. 0 value
                      will ignore this embedding
class_docstr_tok: ["Language","model","embeddings",".","Arguments",":","hidden_size",":","hidden","size","vocab_size",":","vocabulary","size","max_sequence_length",":","maximum","size","of","sequence",".","This","is","used","for","positional","embedding","embedding_dropout_prob",":","dropout","probability","for","embeddings","init_method",":","weight","initialization","method","num_tokentypes",":","size","of","the","token-type","embeddings",".","0","value","will","ignore","this","embedding"]
docstr: Customized load.
docstr_tok: ["Customized","load","."]
code: def load_state_dict(self, state_dict, strict=True):
      """Customized load."""
  
      # Word embedding.
      if self._word_embeddings_key in state_dict:
          state_dict_ = state_dict[self._word_embeddings_key]
      else:
          # for backward compatibility.
          state_dict_ = {}
          for key in state_dict.keys():
              if "word_embeddings" in key:
                  state_dict_[
                      key.split("word_embeddings.")[1]
                  ] = state_dict[key]
      self.word_embeddings.load_state_dict(state_dict_, strict=strict)
  
      # Position embedding.
      if self._position_embeddings_key in state_dict:
          state_dict_ = state_dict[self._position_embeddings_key]
      else:
          # for backward compatibility.
          state_dict_ = {}
          for key in state_dict.keys():
              if "position_embeddings" in key:
                  state_dict_[
                      key.split("position_embeddings.")[1]
                  ] = state_dict[key]
      self.position_embeddings.load_state_dict(
          state_dict_, strict=strict
      )
  
      # Tokentype embedding.
      if self.num_tokentypes > 0:
          state_dict_ = {}
          if self._tokentype_embeddings_key in state_dict:
              state_dict_ = state_dict[self._tokentype_embeddings_key]
          else:
              # for backward compatibility.
              for key in state_dict.keys():
                  if "tokentype_embeddings" in key:
                      state_dict_[
                          key.split("tokentype_embeddings.")[1]
                      ] = state_dict[key]
          if len(state_dict_.keys()) > 0:
              self.tokentype_embeddings.load_state_dict(
                  state_dict_, strict=strict
              )
          else:
              print(
                  "***WARNING*** expected tokentype embeddings in the "
                  "checkpoint but could not find it",
                  flush=True,
              )
  
imports: import torch
  import torch.nn.init
  from torch import Tensor
  from torch import nn
  from torch.nn import functional
  from torch.nn.parameter import Parameter
  from colossalai.legacy.context import ParallelMode, seed
  from colossalai.legacy.core import global_context
  from colossalai.legacy.nn.layer.base_layer import ParallelLayer
  from colossalai.legacy.nn.layer.parallel_1d._utils import gather_forward_split_backward, reduce_grad, reduce_input
  from colossalai.legacy.nn.layer.parallel_1d.layers import Linear1D_Row
  from colossalai.legacy.nn.layer.utils import divide
  from colossalai.legacy.registry import LAYERS, LOSSES
  from colossalai.utils import get_current_device
label_desc: Use image node_id 6 for calling the HiddenParallelEmbedding obj's underlying member method code with example usage: obj.load_state_dict(state_dict, strict) without return types