code: def get_optimizer(args, model):
      optimizer = bmt.optim.AdamOffloadOptimizer(
          model.parameters(), weight_decay=args.weight_decay
      )
      if args.load is not None:
          if os.path.exists(
              os.path.join(
                  args.save, args.save_name + (".rank-%d.opt" % 0)
              )
          ):
              # optimizer state exists
              states = torch.load(
                  os.path.join(
                      args.save,
                      args.save_name + (".rank-%d.opt" % bmt.rank()),
                  )
              )
              optimizer.load_state_dict(states)
      return optimizer
  
imports: import json
  import time
  from typing import Any, Dict, List, Union
  import torch
  import bmtrain
  import os
  from cpm_live.arguments import get_args
  from cpm_live.models import CPMBee, CPMBeeConfig
  from cpm_live.tokenizers import CPMBeeTokenizer
  from cpm_live.utils import allgather_objects, LogManager
  from cpm_live.training_tasks.bee import MixedDataset
label_desc: Use image node_id 3 for calling a global function with example usage: get_optimizer(args, model) and returns: optimizer