code: def safe_set_local_optimizer_state(param, value, optim_state_key):
      """Update the fp32 optimizer state of a partitioned parameter.
      Args:
          param (``torch.nn.Parameter``): A model parameter
          value (``torch.Tensor``): New value
          optim_state_key (``string``): Key value of optimizer state (e.g., `exp_avg` in Adam optimizer)
      """
      # ZeRO stage 3 param
      if hasattr(param, "ds_id"):
          param._z3_optimizer.set_local_hp_param(
              value, param, optim_state_key
          )
  
imports: import torch
  from dataclasses import dataclass
  from deepspeed import comm
  from typing import Dict
label_desc: Use image node_id 12 for calling a global function with example usage: safe_set_local_optimizer_state(param, value, optim_state_key) without return types