class_docstr: Core trainable object in OpenNMT. Implements a trainable interface
  for a simple, generic encoder / decoder or decoder only model.
  
  Args:
    encoder (onmt.encoders.EncoderBase): an encoder object
    decoder (onmt.decoders.DecoderBase): a decoder object
class_docstr_tok: ["Core","trainable","object","in","OpenNMT",".","Implements","a","trainable","interface","for","a","simple",",","generic","encoder","\/","decoder","or","decoder","only","model",".","Args",":","encoder","(","onmt.encoders.EncoderBase",")",":","an","encoder","object","decoder","(","onmt.decoders.DecoderBase",")",":","a","decoder","object"]
docstr: Custom state_dict loading to enable moving module on device as they are loaded
  
  Args:
      model_path: Model path
      precision: same as above
      device: same as above
      strict: same as above
docstr_tok: ["Custom","state_dict","loading","to","enable","moving","module","on","device","as","they","are","loaded","Args",":","model_path",":","Model","path","precision",":","same","as","above","device",":","same","as","above","strict",":","same","as","above"]
code: def load_safe_state_dict(
      self,
      model_path,
      precision=torch.float32,
      device=torch.device("cpu"),
      strict=True,
      offset=0,
  ):
      """Custom state_dict loading to enable moving module on device as they are loaded
  
      Args:
          model_path: Model path
          precision: same as above
          device: same as above
          strict: same as above
      """
      # bitsandbytes quantize weights when .cuda() is called
      # for huge models we need to save Ram
      # so we load the weights  module by module and transfer them to GPU for quantization
      try:
          import safetensors
      except ImportError:
          raise ImportError(
              "run: pip install safetensors, to use safetensors"
          )
      keyfound = {}
      shards = glob(model_path + ".*.safetensors")
      if len(shards) == 0:
          raise ValueError("No safetensors file found")
      f = []
      keys_shard = {}
      for i, shard in enumerate(shards):
          f.append(
              safetensors.safe_open(shard, framework="pt", device="cpu")
          )
          for key in f[i].keys():
              keys_shard[key] = i
      if device == torch.device("cpu"):
          offset = 0
      buf_list = []
      for buf_name, buf in self.named_buffers():
          buf_list.append(buf_name)
      for name, module in self.named_modules():
          named_buf_and_param = list(module.named_buffers()) + list(
              module.named_parameters()
          )
          for param_name, param in named_buf_and_param:
              if len(param_name.split(".")) == 1:  # only last key
                  if name + "." + param_name in keys_shard.keys():
                      ckpt_t = f[
                          keys_shard[name + "." + param_name]
                      ].get_tensor(name + "." + param_name)
                      self._load_param(
                          name,
                          module,
                          param_name,
                          param,
                          buf_list,
                          ckpt_t,
                          offset,
                      )
                      keyfound[name + "." + param_name] = True
                  elif strict and "lora" not in param_name:
                      raise ValueError(
                          "Missing key in safetensors checkpoint: %s"
                          % name
                          + "."
                          + param_name
                      )
                  if precision == torch.int8:
                      torch.quantization.quantize_dynamic(
                          module, inplace=True
                      )
                  else:
                      module.to(precision)
                  module.to(device)
      for key in keys_shard.keys():
          if key not in keyfound.keys() and key not in buf_list:
              raise ValueError(
                  "Extra keys in model state_dict do not match the model config %s"
                  % key
              )
  
imports: import torch
  import torch.nn
  from glob import glob
label_desc: Use image node_id 7 for calling the BaseModel obj's underlying member method code with example usage: obj.load_safe_state_dict(model_path, precision, device, strict, offset) without return types