code: def reset_optimizer(optimizer: torch.optim.Optimizer):
      for group in optimizer.param_groups:
          for param in group["params"]:
              param_state = optimizer.state[param]
              for key in param_state:
                  if "qmap" in key:
                      continue
  
                  if key == "step" and isinstance(
                      param_state[key], int
                  ):
                      param_state[key] = 0
                  else:
                      param_state[key] = torch.zeros_like(
                          param_state[key]
                      )
  
imports: import glob
  import json
  import logging
  import os.path
  import shutil
  from pathlib import Path
  from typing import Dict, List, Sequence
  import bitsandbytes
  import peft
  import safetensors.torch
  import torch
  from huggingface_hub import snapshot_download
  from torch.optim.lr_scheduler import LRScheduler
  from torch.optim.optimizer import Optimizer
  from transformers import TrainerCallback, TrainerControl, TrainerState, TrainingArguments
  from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR
  from axolotl.utils.dict import DictDefault
  from axolotl.utils.distributed import is_main_process
label_desc: Use image node_id 1 for calling a global function with example usage: reset_optimizer(optimizer) without return types