code: def states_option(
      step_num, trainer, ckpt_dir, local_rank=0, option="Saving"
  ):
      """Save or load the trainer states, marked by step_num and local rank."""
      state_path = os.path.join(
          ckpt_dir,
          "{}.states.{}".format(
              str(step_num).zfill(7), str(local_rank).zfill(2)
          ),
      )
      logging.info(
          "[step {}], {} trainer states to/from {}.".format(
              step_num, option, state_path
          )
      )
      if option == "Saving":
          trainer.save_states(state_path)
      elif option == "Loading":
          trainer.load_states(state_path)
      else:
          raise NotImplementedError("Unknown Option: {}".format(option))
  
file_path: gluon-nlp/scripts/pretraining/bert/run_pretraining.py
imports: import os
  import time
  import shutil
  import logging
  import argparse
  import functools
  import collections
  import mxnet
  import numpy
  from pretraining_utils import get_pretrain_data_npz, get_pretrain_data_text, MaskedAccuracy
  from gluonnlp.utils.misc import repeat, grouper, set_seed, init_comm, logging_config, naming_convention
  from gluonnlp.initializer import TruncNorm
  from gluonnlp.models.bert import BertModel, BertForPretrain, get_pretrained_bert
  from gluonnlp.utils.parameter import clip_grad_global_norm
label_desc: Use image node_id 5 for calling a global function with example usage: states_option(step_num, trainer, ckpt_dir, local_rank, option) without return types