#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# Copyright (c) Megvii, Inc. and its affiliates.

import argparse

import megengine as mge
import numpy as np
from megengine import jit

from ..build import build_and_load


def make_parser():
    parser = argparse.ArgumentParser("YOLOX Demo Dump")
    parser.add_argument("-n", "--name", type=str, default="yolox-s", help="model name")
    parser.add_argument("-c", "--ckpt", default=None, type=str, help="ckpt for eval")
    parser.add_argument(
        "--dump_path", default="model.mge", help="path to save the dumped model"
    )
    return parser


def dump_static_graph(model, graph_name="model.mge"):
    model.eval()
    model.head.decode_in_inference = False

    data = mge.Tensor(np.random.random((1, 3, 640, 640)))

    @
jit.trace(capture_as_const=True)