code: def wrap_load_from_state_dict(f):
      @functools.wraps(f)
      def wrapper(
          model,
          state_dict,
          prefix,
          local_metadata,
          strict,
          missing_keys,
          unexpected_keys,
          error_msgs,
      ):
          replaced_old_value = None
          key = None
          # We may have nested wrappers if we launch multiple initialization context.
          # Use state_dict_quantized flag to quantize state_dict only once
          if (
              hasattr(model.weight, "weight_quantized")
              and getattr(model.weight, "weight_quantized")
              and not hasattr(model.weight, "state_dict_quantized")
          ):
              setattr(model.weight, "state_dict_quantized", True)
              key = prefix + "weight"
              if key in state_dict:
                  (
                      quantized_weight,
                      quant_scale,
                      quant_min,
                  ) = model.weight.quantizer.quantize(state_dict[key])
                  quantized_weight = quantized_weight.view(
                      model.weight.dtype
                  )
                  quant_scale = quant_scale.view(model.weight.dtype)
                  quant_min = quant_min.view(model.weight.dtype)
  
                  replaced_old_value = state_dict[key]
  
                  state_dict[key] = concat_to_compat_param(
                      quantized_weight, quant_scale, quant_min
                  )
  
          f(
              model,
              state_dict,
              prefix,
              local_metadata,
              strict,
              missing_keys,
              unexpected_keys,
              error_msgs,
          )
  
          if replaced_old_value is not None:
              state_dict[key] = replaced_old_value
              delattr(model.weight, "state_dict_quantized")
  
      return wrapper
  
imports: import torch
  import deepspeed
  from torch import Tensor
  from typing import Tuple
  import torch.nn
  from typing import Dict, Callable, Union
  from deepspeed.accelerator import get_accelerator
  import functools
label_desc: Use image node_id 9 for calling a global function with example usage: wrap_load_from_state_dict(f) and returns: wrapper