code: def save_model(
      output_path, model, steps, opt, lr_scheduler, max_ckpt=2
  ):
      if paddle.distributed.get_rank() == 0:
          output_dir = os.path.join(output_path, "model_%d" % steps)
          if not os.path.exists(output_dir):
              os.makedirs(output_dir)
  
          paddle.save(
              model.state_dict(),
              os.path.join(output_dir, "ckpt.pdparams"),
          )
          #         paddle.save({ "global_step": steps }, os.path.join(output_dir, "step"))
          #         paddle.save(opt.state_dict(), os.path.join(output_dir, "opt.pdparams"))
          #         paddle.save(lr_scheduler.state_dict(), os.path.join(output_dir, "lr_scheduler.pdparams"))
  
          log.info("save model %s" % output_dir)
  
          ckpt_paths = glob.glob(os.path.join(output_path, "model_*"))
          if len(ckpt_paths) > max_ckpt:
  
              def version(x):
                  x = int(x.split("_")[-1])
                  return x
  
              rm_ckpt_paths = sorted(
                  ckpt_paths, key=version, reverse=True
              )[max_ckpt:]
              for ckpt_dir in rm_ckpt_paths:
                  if os.path.exists(ckpt_dir):
                      shutil.rmtree(ckpt_dir)
  
file_path: PGL/ogb_examples/lsc/MAG240M/utils.py
imports: import os
  import paddle
  import paddle.distributed
  import glob
  import shutil
  from pgl.utils.logger import log
label_desc: Use image node_id 2 for calling a global function with example usage: save_model(output_path, model, steps, opt, lr_scheduler, max_ckpt) and returns: x