from arch import resnet_nas, vgg_atl

def check_model(opt):
    if opt.model.startswith('resnet'):
        model = resnet_nas.__dict__[opt.model](pretrained=opt.pretrained_model)
        return model
    elif opt.model.startswith('vgg'):
        model = vgg_atl.__dict__[opt.model](pretrained=opt.pretrained_model)
        return model
    else:
        raise Exception('Unknown model')
