import os
import argparse



parser = argparse.ArgumentParser()

parser.add_argument('-a','--arch', required=True, help='a network name which will be trained')
parser.add_argument('--path', required=True, help='path of the config file')


"""
You should add below lines to model_zoo.py in order to train the network.
"""

if __name__ == '__main__':
    args = parser.parse_args()
    with open('proxylessnas/training/net224x224/proxyless_nas/model_zoo.py', 'a') as f:
        arch = f"""
{args.arch} = partial(
    proxyless_base,
    pretrained=False,
    net_config=\'{args.path}\',
    net_weight=None
)

"""
        f.write(arch)
