def METHOD_NAME(args):
    paddle.set_device("gpu" if args.n_gpu else "cpu")
    args.model_type = args.model_type.lower()
    args.task_name = args.task_name.lower()
    model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
    config_path = os.path.join(args.model_name_or_path, "config.json")
    cfg_dict = dict(json.loads(open(config_path).read()))
    kept_layers_index = {}
    if args.depth_mult < 1.0:
        depth = round(cfg_dict["init_args"][0]["num_hidden_layers"] * args.depth_mult)
        cfg_dict["init_args"][0]["num_hidden_layers"] = depth
        for idx, i in enumerate(range(1, depth + 1)):
            kept_layers_index[idx] = math.floor(i / args.depth_mult) - 1
    os.rename(config_path, config_path + "_bak")
    with open(config_path, "w", encoding="utf-8") as f:
        f.write(json.dumps(cfg_dict, ensure_ascii=False))
    model = model_class.from_pretrained(args.model_name_or_path)
    origin_model = model_class.from_pretrained(args.model_name_or_path)
    os.rename(config_path + "_bak", config_path)
    sp_config = supernet(expand_ratio=[1.0, args.width_mult])
    model = Convert(sp_config).convert(model)
    ofa_model = OFA(model)
    sd = paddle.load(os.path.join(args.model_name_or_path, "model_state.pdparams"))
    if len(kept_layers_index) == 0:
        ofa_model.model.set_state_dict(sd)
    else:
        for name, params in ofa_model.model.named_parameters():
            if "encoder" not in name:
                params.set_value(sd[name])
            else:
                idx = int(name.strip().split(".")[3])
                mapping_name = name.replace("." + str(idx) + ".", "." + str(kept_layers_index[idx]) + ".")
                params.set_value(sd[mapping_name])
    best_config = utils.dynabert_config(ofa_model, args.width_mult)
    for name, sublayer in ofa_model.model.named_sublayers():
        if isinstance(sublayer, paddle.nn.MultiHeadAttention):
            sublayer.num_heads = int(args.width_mult * sublayer.num_heads)
    origin_model_new = ofa_model.export(
        best_config,
        input_shapes=[[1, args.max_seq_length], [1, args.max_seq_length]],
        input_dtypes=["int64", "int64"],
        origin_model=origin_model,
    )
    for name, sublayer in origin_model_new.named_sublayers():
        if isinstance(sublayer, paddle.nn.MultiHeadAttention):
            sublayer.num_heads = int(args.width_mult * sublayer.num_heads)
    output_dir = os.path.join(args.sub_model_output_dir, "model_width_%.5f" % args.width_mult)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    model_to_save = origin_model_new
    model_to_save.save_pretrained(output_dir)
    input_spec = [
        paddle.static.InputSpec(shape=[None, None], dtype="int64"),  # input_ids
        paddle.static.InputSpec(shape=[None, None], dtype="int64"),  # token_type_ids
    ]
    origin_model_new = paddle.jit.to_static(origin_model_new, input_spec=input_spec)
    paddle.jit.save(origin_model_new, args.static_sub_model)