code: def _load_checkpoint(
      self, load_dir, load_module_strict=True, tag=None
  ):
      is_pipe_parallel = isinstance(self.module, PipelineModule)
      if is_pipe_parallel:
          raise RuntimeError(
              "pipeline parallelism is currently not supported in inference."
          )
      if not isinstance(load_dir, dict) and os.path.isdir(load_dir):
          if tag is None:
              latest_path = os.path.join(load_dir, "latest")
              if os.path.isfile(latest_path):
                  with open(latest_path, "r") as fd:
                      tag = fd.read().strip()
  
          ckpt_list = self._get_all_ckpt_names(load_dir, tag)
          sd_loader = SDLoaderFactory.get_sd_loader(
              ckpt_list, self.checkpoint_engine
          )
      else:
          sd_loader = SDLoaderFactory.get_sd_loader_json(
              load_dir, self.checkpoint_engine
          )
  
      checkpoint = sd_loader["checkpoints"]
  
      if type(checkpoint) is list:
          self.sd = torch.load(checkpoint[0], map_location="cpu")
          self.key_list = list(self.sd.keys())
  
          self.load_model_with_checkpoint(self.module)
  
          for i in range(1, len(checkpoint)):
              if not dist.is_initialized() or dist.get_rank() == 0:
                  print(f"loading checkpoint ({i})")
              self.sd = torch.load(
                  checkpoint[i],
                  map_location=get_accelerator().device_name(),
              )
              self.key_list = list(self.sd.keys())
              self.load_model_with_checkpoint(self.module)
      else:
          mp_rank = (
              0
              if self.mpu is None
              else self.mpu.get_model_parallel_rank()
          )
  
          load_path, checkpoint, quantize_config = sd_loader.load(
              self._config.tensor_parallel.tp_size,
              mp_rank,
              is_pipe_parallel=is_pipe_parallel,
              quantize=(self._config.dtype is torch.int8),
              quantize_groups=self.quantize_groups,
              mlp_extra_grouping=self.mlp_extra_grouping,
          )
  
          (
              self.quantization_scales,
              self.quantize_merge_count,
          ) = quantize_config
  
          moe, _ = has_moe_layers(self.module)
          if moe:
              from deepspeed.runtime.engine import DeepSpeedEngine
  
              old_moe_load = False
              if not isinstance(checkpoint["num_experts"], list):
                  old_moe_load = True
              DeepSpeedEngine.load_moe_state_dict(
                  load_dir,
                  tag,
                  state_dict=checkpoint[
                      self._choose_module_key(checkpoint)
                  ],
                  old_moe_load=old_moe_load,
                  model=self.module,
                  mpu=self.mpu,
                  checkpoint_engine=self.checkpoint_engine,
              )
  
          self.module.load_state_dict(
              state_dict=checkpoint[
                  self._choose_module_key(checkpoint)
              ],
              strict=load_module_strict,
          )
  
imports: import torch
  import time
  import os
  from deepspeed import comm
  from deepspeed.utils.logging import log_dist
  from torch.nn.modules import Module
  from packaging import version
  from deepspeed.runtime.checkpoint_engine.torch_checkpoint_engine import TorchCheckpointEngine
  from deepspeed.utils.timer import SynchronizedWallClockTimer
  from ..runtime.state_dict_factory import SDLoaderFactory
  from ..runtime.weight_quantizer import WeightQuantization
  from ..module_inject import replace_transformer_layer, generic_injection
  from ..comm.comm import init_distributed
  from ..pipe import PipelineModule
  from ..moe.utils import has_moe_layers
  from ..module_inject import LinearAllreduce, LinearLayer, Normalize, ReplaceWithTensorSlicing
  from deepspeed.accelerator import get_accelerator
  from ..module_inject.policy import TransformerPolicy
  from ..module_inject.auto_tp import AutoTP
  from ..module_inject.replace_policy import generic_policies
  from ..module_inject.auto_tp_model_utils import build_bloom_alibi_tensor, build_mpt_atten_bias_tensor, build_mpt_alibi_tensor, get_alibi_mask
  from ..ops.transformer.inference.ds_attention import DeepSpeedSelfAttention
  from ..model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference
  from torch import nn
label_desc: Use image node_id 18 for calling the InferenceEngine obj's underlying member method code with example usage: obj._load_checkpoint(load_dir, load_module_strict, tag) without return types