from uci import train_uci

Datasets = ['boston', 'concrete', 'energy', 'wine', 'yacht']
Depths = [2, 3, 4, 5]
splits = 20

for dataset in Datasets:
    if dataset == 'protein':
        splits = 5
    else:
        splits = 20
    for depth in Depths:
        for split in range(splits):
            for model in ['dgp', 'dwp', 'idwp', 'ibdwp']:
                output_fn = 'improved_uci_results/{}_{}_d{}_s{}'.format(model, depth, dataset, split)
                train_uci(output_fn, model, dataset, split, depth, device='cpu')

