def METHOD_NAME(args):
    os.environ['PADDLESEG_EXPORT_STAGE'] = 'True'
    cfg = Config(args.config)
    net = cfg.model
    if args.model_path is not None:
        para_state_dict = paddle.load(args.model_path)
        net.set_dict(para_state_dict)
        logger.info('Loaded trained params of model successfully.')
    if args.input_shape is None:
        shape = [None, 3, None, None]
    else:
        shape = args.input_shape
    output_op = args.output_op
    if args.without_argmax:
        logger.warning(
            '--without_argmax will be deprecated, please use --output_op')
        output_op = 'none'
    if args.with_softmax:
        logger.warning(
            '--with_softmax will be deprecated, please use --output_op')
        output_op = 'softmax'
    new_net = net if output_op == 'none' else SavedSegmentationNet(net,
                                                                   output_op)
    new_net.eval()
    new_net = paddle.jit.to_static(
        new_net,
        input_spec=[paddle.static.InputSpec(
            shape=shape, dtype='float32')])
    save_path = os.path.join(args.save_dir, 'model')
    paddle.jit.save(new_net, save_path)
    yml_file = os.path.join(args.save_dir, 'deploy.yaml')
    with open(yml_file, 'w') as file:
        transforms = cfg.export_config.get('transforms', [{
            'type': 'Normalize'
        }])
        output_dtype = 'int32' if output_op == 'argmax' else 'float32'
        data = {
            'Deploy': {
                'model': 'model.pdmodel',
                'params': 'model.pdiparams',
                'transforms': transforms,
                'input_shape': shape,
                'output_op': output_op,
                'output_dtype': output_dtype
            }
        }
        yaml.dump(data, file)
    logger.info(f'The inference model is saved in {args.save_dir}')