code: def save(self, program, lr_scheduler=None, epoch=0, for_train=True):
      model_save_dir = os.path.join(self.model_save_dir, str(epoch))
      if not os.path.exists(model_save_dir):
          # may be more than one processes trying
          # to create the directory
          try:
              os.makedirs(model_save_dir)
          except OSError as exc:
              if exc.errno != errno.EEXIST:
                  raise
              pass
  
      param_state_dict = program.state_dict(mode="param")
      for name, param in param_state_dict.items():
          # for non dist param, we only save their at rank 0,
          # but for dist param, we need to save their at all ranks.
          if "dist@" in name and "@rank@" in name or self.rank == 0:
              paddle.save(
                  param, os.path.join(model_save_dir, name + ".pdparam")
              )
  
      if for_train:
          opt_state_dict = program.state_dict(mode="opt")
          for name, opt in opt_state_dict.items():
              if "@GRAD" in name:
                  continue
              # for non dist opt var, we only save their at rank 0,
              # but for dist opt var, we need to save their at all ranks.
              if "dist@" in name and "@rank@" in name or self.rank == 0:
                  paddle.save(
                      opt, os.path.join(model_save_dir, name + ".pdopt")
                  )
  
          if self.rank == 0:
              # save some extra info for resume
              # pretrain_world_size, embedding_size, num_classes are used for
              # re-split fc weight when gpu setting changed.
              # epoch use to restart.
              config_file = os.path.join(model_save_dir, "meta.json")
              extra_info = dict()
              extra_info["pretrain_world_size"] = self.world_size
              extra_info["embedding_size"] = self.embedding_size
              extra_info["num_classes"] = self.num_classes
              extra_info["epoch"] = epoch
              extra_info["lr_state"] = lr_scheduler.state_dict()
              with open(config_file, "w") as f:
                  json.dump(extra_info, f)
  
      logging.info("Save model to {}.".format(model_save_dir))
      if self.rank == 0 and self.max_num_last_checkpoint > 0:
          for idx in range(
              -1, epoch - self.max_num_last_checkpoint + 1
          ):
              path = os.path.join(self.model_save_dir, str(idx))
              if os.path.exists(path):
                  logging.info("Remove checkpoint {}.".format(path))
                  shutil.rmtree(path)
  
file_path: insightface/recognition/arcface_paddle/static/utils/io.py
imports: import errno
  import os
  import paddle
  import logging
  import numpy
  import shutil
  import json
  from utils.rearrange_weight import rearrange_weight
label_desc: Use image node_id 2 for calling the Checkpoint obj's underlying member method code with example usage: obj.save(program, lr_scheduler, epoch, for_train) without return types