from .logistic import Logistic
from .lenet import LeNet
from .charlstm import CharLSTM

def gen_model(dataset):
    if dataset == 'synthetic':
        w = Logistic(60, 10)
    elif dataset == 'cifar10':
        w = LeNet()
    elif dataset == 'shakespeare':
        w = CharLSTM()
    return w