code: def main(args):
      if (
          args.begin is not None
          or args.end is not None
          or args.ids is not None
      ):
          print(
              f"Before filtering, the checkpoints are {args.checkpoints}"
          )
          prefix = os.path.commonprefix(args.checkpoints)
          postfix = os.path.commonprefix(
              [ele[::-1] for ele in args.checkpoints]
          )[::-1]
          checkpoint_id_l = [
              int(ele[len(prefix) : -len(postfix)])
              for ele in args.checkpoints
          ]
          ckpt_paths = []
          if args.ids is not None:
              for ele in args.ids:
                  assert ele in checkpoint_id_l
                  ckpt_paths.append(f"{prefix}{ele}{postfix}")
          else:
              assert (
                  args.begin is not None and args.end is not None
              ), "Must specify both begin and end if you want to select a range!"
              assert args.begin >= 0
              assert args.end >= args.begin
              for ele in checkpoint_id_l:
                  if ele >= args.begin and ele <= args.end:
                      ckpt_paths.append(f"{prefix}{ele}{postfix}")
      else:
          ckpt_paths = args.checkpoints
      print(f"Load models from {ckpt_paths}")
      print(
          "Average the models and save it to {}".format(args.save_path)
      )
      assert (
          len(ckpt_paths) > 0
      ), "Cannot found checkpoints. You may need to check the inputs again."
      res = mx.npx.load(ckpt_paths[0])
      keys = res.keys()
      for ckpt_path in ckpt_paths[1:]:
          ckpt = mx.npx.load(ckpt_path)
          for key in keys:
              res[key] += ckpt[key]
      for key in keys:
          res[key] /= len(ckpt_paths)
      mx.npx.savez(args.save_path, **res)
  
imports: import argparse
  import mxnet
  import os
label_desc: Use image node_id 2 for calling a global function with example usage: main(args) without return types