--- a
+++ b/eval.py
@@ -0,0 +1,590 @@
+
+import yaml
+import tensorboard
+import torch
+import os
+import shutil
+import sys
+import csv
+import argparse
+import pickle
+from models.resnet_simclr import ResNetSimCLR
+from clinical_ts.cpc import CPCModel
+import torch.nn.functional as F
+from tqdm import tqdm
+import numpy as np
+import matplotlib.pyplot as plt
+from sklearn.decomposition import PCA
+from sklearn.manifold import TSNE
+from sklearn.metrics import roc_auc_score
+from clinical_ts.simclr_dataset_wrapper import SimCLRDataSetWrapper
+from clinical_ts.eval_utils_cafa import eval_scores, eval_scores_bootstrap
+from clinical_ts.timeseries_utils import aggregate_predictions
+import pdb
+from copy import deepcopy
+from os.path import join, isdir
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+
+def parse_args():
+    parser = argparse.ArgumentParser("Finetuning tests")
+    parser.add_argument("--model_file")
+    parser.add_argument("--method")
+    parser.add_argument("--dataset", nargs="+", default="./data/ptb_xl_fs100")
+    parser.add_argument("--batch_size", type=int, default=512)
+    parser.add_argument("--discriminative_lr", default=False, action="store_true")
+    parser.add_argument("--num_workers", type=int, default=8)
+    parser.add_argument("--hidden", default=False, action="store_true")
+    parser.add_argument("--lr_schedule", default="{}")
+    parser.add_argument("--use_pretrained", default=False, action="store_true")
+    parser.add_argument("--linear_evaluation",
+                        default=False, action="store_true", help="use linear evaluation")
+    parser.add_argument("--test_noised", default=False, action="store_true", help="validate also on a noisy dataset")
+    parser.add_argument("--noise_level", default=1, type=int, help="level of noise induced to the second validations set")
+    parser.add_argument("--folds", default=8, type=int, help="number of folds used in finetuning (between 1-8)")
+    parser.add_argument("--tag", default="")
+    parser.add_argument("--eval_only", action="store_true", default=False, help="only evaluate mode")
+    parser.add_argument("--load_finetuned", action="store_true", default=False)
+    parser.add_argument("--test", action="store_true", default=False)
+    parser.add_argument("--verbose", action="store_true", default=False)
+    parser.add_argument("--cpc", action="store_true", default=False)
+    parser.add_argument("--model_location")
+    parser.add_argument("--l_epochs", type=int, default=0, help="number of head-only epochs (these are performed first)")
+    parser.add_argument("--f_epochs", type=int, default=0, help="number of finetuning epochs (these are perfomed after head-only training")
+    parser.add_argument("--normalize", action="store_true", default=False, help="normalize dataset with ptbxl mean and std")
+    parser.add_argument("--bn_head", action="store_true", default=False)
+    parser.add_argument("--ps_head", type=float, default=0.0)
+    parser.add_argument("--conv_encoder", action="store_true", default=False)
+    parser.add_argument("--base_model", default="xresnet1d50")
+    parser.add_argument("--widen", default=1, type=int, help="use wide xresnet1d50")
+    args = parser.parse_args()
+    return args
+
+
+def get_new_state_dict(init_state_dict, lightning_state_dict, method="simclr"):
+    # in case of moco model
+    from collections import OrderedDict
+    # lightning_state_dict = lightning_state_dict["state_dict"]
+    new_state_dict = OrderedDict()
+    if method != "cpc":
+        if method == "moco":
+            for key in init_state_dict:
+                l_key = "encoder_q." + key
+                if l_key in lightning_state_dict.keys():
+                    new_state_dict[key] = lightning_state_dict[l_key]
+        elif method == "simclr":
+            for key in init_state_dict:
+                if "features" in key:
+                    l_key = key.replace("features", "encoder.features")
+                if l_key in lightning_state_dict.keys():
+                    new_state_dict[key] = lightning_state_dict[l_key]
+        elif method == "swav":
+
+            for key in init_state_dict:
+                if "features" in key:
+                    l_key = key.replace("features", "model.features")
+                if l_key in lightning_state_dict.keys():
+                    new_state_dict[key] = lightning_state_dict[l_key]
+        elif method == "byol":
+            for key in init_state_dict:
+                l_key = "online_network.encoder." + key
+                if l_key in lightning_state_dict.keys():
+                    new_state_dict[key] = lightning_state_dict[l_key]
+        else:
+            raise("method unknown")
+        new_state_dict["l1.weight"] = init_state_dict["l1.weight"]
+        new_state_dict["l1.bias"] = init_state_dict["l1.bias"]
+        if "l2.weight" in init_state_dict.keys():
+            new_state_dict["l2.weight"] = init_state_dict["l2.weight"]
+            new_state_dict["l2.bias"] = init_state_dict["l2.bias"]
+
+        assert(len(init_state_dict) == len(new_state_dict))
+    else:
+        for key in init_state_dict:
+            l_key = "model_cpc." + key
+            if l_key in lightning_state_dict.keys():
+                new_state_dict[key] = lightning_state_dict[l_key]
+            if "head" in key:
+                new_state_dict[key] = init_state_dict[key]
+    return new_state_dict
+
+
+def adjust(model, num_classes, hidden=False):
+    in_features = model.l1.in_features
+    last_layer = torch.nn.modules.linear.Linear(
+        in_features, num_classes).to(device)
+    if hidden:
+        model.l1 = torch.nn.modules.linear.Linear(
+            in_features, in_features).to(device)
+        model.l2 = last_layer
+    else:
+        model.l1 = last_layer
+
+    def def_forward(self):
+        def new_forward(x):
+            h = self.features(x)
+            h = h.squeeze()
+
+            x = self.l1(h)
+            if hidden:
+                x = F.relu(x)
+                x = self.l2(x)
+            return x
+        return new_forward
+
+    model.forward = def_forward(model)
+
+
+def configure_optimizer(model, batch_size, head_only=False, discriminative_lr=False, base_model="xresnet1d", optimizer="adam", discriminative_lr_factor=1):
+    loss_fn = F.binary_cross_entropy_with_logits
+    if base_model == "xresnet1d":
+        wd = 1e-1
+        if head_only:
+            lr = (8e-3*(batch_size/256))
+            optimizer = torch.optim.AdamW(
+                model.l1.parameters(), lr=lr, weight_decay=wd)
+        else:
+            lr = 0.01
+            if not discriminative_lr:
+                optimizer = torch.optim.AdamW(
+                    model.parameters(), lr=lr, weight_decay=wd)
+            else:
+                param_dict = dict(model.named_parameters())
+                keys = param_dict.keys()
+                weight_layer_nrs = set()
+                for key in keys:
+                    if "features" in key:
+                        # parameter names have the form features.x
+                        weight_layer_nrs.add(key[9])
+                weight_layer_nrs = sorted(weight_layer_nrs, reverse=True)
+                features_groups = []
+                while len(weight_layer_nrs) > 0:
+                    if len(weight_layer_nrs) > 1:
+                        features_groups.append(list(filter(
+                            lambda x: "features." + weight_layer_nrs[0] in x or "features." + weight_layer_nrs[1] in x,  keys)))
+                        del weight_layer_nrs[:2]
+                    else:
+                        features_groups.append(
+                            list(filter(lambda x: "features." + weight_layer_nrs[0] in x,  keys)))
+                        del weight_layer_nrs[0]
+                # filter linear layers
+                linears = list(filter(lambda x: "l" in x, keys))
+                groups = [linears] + features_groups
+                optimizer_param_list = []
+                tmp_lr = lr
+
+                for layers in groups:
+                    layer_params = [param_dict[param_name]
+                                    for param_name in layers]
+                    optimizer_param_list.append(
+                        {"params": layer_params, "lr": tmp_lr})
+                    tmp_lr /= 4
+                optimizer = torch.optim.AdamW(
+                    optimizer_param_list, lr=lr, weight_decay=wd)
+
+        print("lr", lr)
+        print("wd", wd)
+        print("batch size", batch_size)
+
+    elif base_model == "cpc":
+        if(optimizer == "sgd"):
+            opt = torch.optim.SGD
+        elif(optimizer == "adam"):
+            opt = torch.optim.AdamW
+        else:
+            raise NotImplementedError("Unknown Optimizer.")
+        lr = 1e-4
+        wd = 1e-3
+        if(head_only):
+            lr = 1e-3
+            print("Linear eval: model head", model.head)
+            optimizer = opt(model.head.parameters(), lr, weight_decay=wd)
+        elif(discriminative_lr_factor != 1.):  # discrimative lrs
+            optimizer = opt([{"params": model.encoder.parameters(), "lr": lr*discriminative_lr_factor*discriminative_lr_factor}, {
+                            "params": model.rnn.parameters(), "lr": lr*discriminative_lr_factor}, {"params": model.head.parameters(), "lr": lr}], lr, weight_decay=wd)
+            print("Finetuning: model head", model.head)
+            print("discriminative lr: ", discriminative_lr_factor)
+        else:
+            lr = 1e-3
+            print("normal supervised training")
+            optimizer = opt(model.parameters(), lr, weight_decay=wd)
+    else:
+        raise("model unknown")
+    return loss_fn, optimizer
+
+
+def load_model(linear_evaluation, num_classes, use_pretrained, discriminative_lr=False, hidden=False, conv_encoder=False, bn_head=False, ps_head=0.5, location="./checkpoints/moco_baselinewonder200.ckpt", method="simclr", base_model="xresnet1d50", out_dim=16, widen=1):
+    discriminative_lr_factor = 1
+    if use_pretrained:
+        print("load model from " + location)
+        discriminative_lr_factor = 0.1
+        if base_model == "cpc":
+            lightning_state_dict = torch.load(location, map_location=device)
+
+            # num_head = np.sum([1 if 'proj' in f else 0 for f in lightning_state_dict.keys()])
+            if linear_evaluation:
+                lin_ftrs_head = []
+                bn_head = False
+                ps_head = 0.0
+            else:
+                if hidden:
+                    lin_ftrs_head = [512]
+                else:
+                    lin_ftrs_head = []
+
+            if conv_encoder:
+                strides = [2, 2, 2, 2]
+                kss = [10, 4, 4, 4]
+            else:
+                strides = [1]*4
+                kss = [1]*4
+
+            model = CPCModel(input_channels=12, strides=strides, kss=kss, features=[512]*4, n_hidden=512, n_layers=2, mlp=False, lstm=True, bias_proj=False,
+                             num_classes=num_classes, skip_encoder=False, bn_encoder=True, lin_ftrs_head=lin_ftrs_head, ps_head=ps_head, bn_head=bn_head).to(device)
+
+            if "state_dict" in lightning_state_dict.keys():
+                print("load pretrained model")
+                model_state_dict = get_new_state_dict(
+                    model.state_dict(), lightning_state_dict["state_dict"], method="cpc")
+            else:
+                print("load already finetuned model")
+                model_state_dict = lightning_state_dict
+            model.load_state_dict(model_state_dict)
+        else:
+            model = ResNetSimCLR(base_model, out_dim, hidden=hidden, widen=widen).to(device)
+            model_state_dict = torch.load(location, map_location=device)
+            if "state_dict" in model_state_dict.keys():
+                model_state_dict = model_state_dict["state_dict"]
+            if "l1.weight" in model_state_dict.keys():  # load already fine-tuned model
+                model_classes = model_state_dict["l1.weight"].shape[0]
+                if model_classes != num_classes:
+                    raise Exception("Loaded model has different output dim ({}) than needed ({})".format(
+                        model_classes, num_classes))
+                adjust(model, num_classes, hidden=hidden)
+                if not hidden and "l2.weight" in model_state_dict:
+                    del model_state_dict["l2.weight"]
+                    del model_state_dict["l2.bias"]
+                model.load_state_dict(model_state_dict)
+            else:  # load pretrained model
+                base_dict = model.state_dict()
+                model_state_dict = get_new_state_dict(
+                    base_dict, model_state_dict, method=method)
+                model.load_state_dict(model_state_dict)
+                adjust(model, num_classes, hidden=hidden)
+
+    else:
+        if "xresnet1d" in base_model:
+            model = ResNetSimCLR(base_model, out_dim, hidden=hidden, widen=widen).to(device)
+            adjust(model, num_classes, hidden=hidden)
+        elif base_model == "cpc":
+            if linear_evaluation:
+                lin_ftrs_head = []
+                bn_head = False
+                ps_head = 0.0
+            else:
+                if hidden:
+                    lin_ftrs_head = [512]
+                else:
+                    lin_ftrs_head = []
+
+            if conv_encoder:
+                strides = [2, 2, 2, 2]
+                kss = [10, 4, 4, 4]
+            else:
+                strides = [1]*4
+                kss = [1]*4
+
+            model = CPCModel(input_channels=12, strides=strides, kss=kss, features=[512]*4, n_hidden=512, n_layers=2, mlp=False, lstm=True, bias_proj=False,
+                             num_classes=num_classes, skip_encoder=False, bn_encoder=True, lin_ftrs_head=lin_ftrs_head, ps_head=ps_head, bn_head=bn_head).to(device)
+
+        else:
+            raise Exception("model unknown")
+
+    return model
+
+
+def evaluate(model, dataloader, idmap, lbl_itos, cpc=False):
+    preds, targs = eval_model(model, dataloader, cpc=cpc)
+    scores = eval_scores(targs, preds, classes=lbl_itos, parallel=True)
+    preds_agg, targs_agg = aggregate_predictions(preds, targs, idmap)
+    scores_agg = eval_scores(targs_agg, preds_agg,
+                             classes=lbl_itos, parallel=True)
+    macro = scores["label_AUC"]["macro"]
+    macro_agg = scores_agg["label_AUC"]["macro"]
+    return preds, macro, macro_agg
+
+
+def set_train_eval(model, cpc, linear_evaluation):
+    if linear_evaluation:
+        if cpc:
+            model.encoder.eval()
+        else:
+            model.features.eval()
+    else:
+        model.train()
+
+
+def train_model(model, train_loader, valid_loader, test_loader, epochs, loss_fn, optimizer, head_only=True, linear_evaluation=False, percentage=1, lr_schedule=None, save_model_at=None, val_idmap=None, test_idmap=None, lbl_itos=None, cpc=False):
+    if head_only:
+        if linear_evaluation:
+            print("linear evaluation for {} epochs".format(epochs))
+        else:
+            print("head-only for {} epochs".format(epochs))
+    else:
+        print("fine tuning for {} epochs".format(epochs))
+
+    if head_only:
+        for key, param in model.named_parameters():
+            if "l1." not in key and "head." not in key:
+                param.requires_grad = False
+        print("copying state dict before training for sanity check after training")
+
+    else:
+        for param in model.parameters():
+            param.requires_grad = True
+    if cpc:
+        data_type = model.encoder[0][0].weight.type()
+    else:
+        data_type = model.features[0][0].weight.type()
+
+    set_train_eval(model, cpc, linear_evaluation)
+    state_dict_pre = deepcopy(model.state_dict())
+    print("epoch", "batch", "loss\n========================")
+    loss_per_epoch = []
+    macro_agg_per_epoch = []
+    max_batches = len(train_loader)
+    break_point = int(percentage*max_batches)
+    best_macro = 0
+    best_macro_agg = 0
+    best_epoch = 0
+    best_preds = None
+    test_macro = 0
+    test_macro_agg = 0
+    for epoch in tqdm(range(epochs)):
+        if type(lr_schedule) == dict:
+            if epoch in lr_schedule.keys():
+                for param_group in optimizer.param_groups:
+                    param_group['lr'] /= lr_schedule[epoch]
+        total_loss_one_epoch = 0
+        for batch_idx, samples in enumerate(train_loader):
+            if batch_idx == break_point:
+                print("break at batch nr.", batch_idx)
+                break
+            data = samples[0].to(device).type(data_type)
+            labels = samples[1].to(device).type(data_type)
+            optimizer.zero_grad()
+            preds = model(data)
+            loss = loss_fn(preds, labels)
+            loss.backward()
+            optimizer.step()
+            total_loss_one_epoch += loss.item()
+            if(batch_idx % 100 == 0):
+                print(epoch, batch_idx, loss.item())
+        loss_per_epoch.append(total_loss_one_epoch)
+
+        preds, macro, macro_agg = evaluate(
+            model, valid_loader, val_idmap, lbl_itos, cpc=cpc)
+        macro_agg_per_epoch.append(macro_agg)
+
+        print("loss:", total_loss_one_epoch)
+        print("aggregated macro:", macro_agg)
+        if macro_agg > best_macro_agg:
+            torch.save(model.state_dict(), save_model_at)
+            best_macro_agg = macro_agg
+            best_macro = macro
+            best_epoch = epoch
+            best_preds = preds
+            _, test_macro, test_macro_agg = evaluate(
+                model, test_loader, test_idmap, lbl_itos, cpc=cpc)
+
+        set_train_eval(model, cpc, linear_evaluation)
+
+    if epochs > 0:
+        sanity_check(model, state_dict_pre, linear_evaluation, head_only)
+    return loss_per_epoch, macro_agg_per_epoch, best_macro, best_macro_agg, test_macro, test_macro_agg, best_epoch, best_preds
+
+
+def sanity_check(model, state_dict_pre, linear_evaluation, head_only):
+    """
+    Linear classifier should not change any weights other than the linear layer.
+    This sanity check asserts nothing wrong happens (e.g., BN stats updated).
+    """
+    print("=> loading state dict for sanity check")
+    state_dict = model.state_dict()
+    if linear_evaluation:
+        for k in list(state_dict.keys()):
+            # only ignore fc layer
+            if 'fc.' in k or 'head.' in k or 'l1.' in k:
+                continue
+
+            equals = (state_dict[k].cpu() == state_dict_pre[k].cpu()).all()
+            if (linear_evaluation != equals):
+                raise Exception(
+                    '=> failed sanity check in {}'.format("linear_evaluation"))
+    elif head_only:
+        for k in list(state_dict.keys()):
+            # only ignore fc layer
+            if 'fc.' in k or 'head.' in k:
+                continue
+
+            equals = (state_dict[k].cpu() == state_dict_pre[k].cpu()).all()
+            if (equals and "running_mean" in k):
+                raise Exception(
+                    '=> failed sanity check in {}'.format("head-only"))
+    # else:
+    #     for k in list(state_dict.keys()):
+    #         equals=(state_dict[k].cpu() == state_dict_pre[k].cpu()).all()
+    #         if equals:
+    #             pdb.set_trace()
+    #             raise Exception('=> failed sanity check in {}'.format("fine_tuning"))
+
+    print("=> sanity check passed.")
+
+
+def eval_model(model, valid_loader, cpc=False):
+    if cpc:
+        data_type = model.encoder[0][0].weight.type()
+    else:
+        data_type = model.features[0][0].weight.type()
+    model.eval()
+    preds = []
+    targs = []
+    with torch.no_grad():
+        for batch_idx, samples in tqdm(enumerate(valid_loader)):
+            data = samples[0].to(device).type(data_type)
+            preds_tmp = torch.sigmoid(model(data))
+            targs.append(samples[1])
+            preds.append(preds_tmp.cpu())
+        preds = torch.cat(preds).numpy()
+        targs = torch.cat(targs).numpy()
+
+    return preds, targs
+
+
+def get_dataset(batch_size, num_workers, target_folder, apply_noise=False, percentage=1.0, folds=8, t_params=None, test=False, normalize=False):
+    if apply_noise:
+        transformations = ["BaselineWander",
+                           "PowerlineNoise", "EMNoise", "BaselineShift"]
+        if normalize:
+            transformations.append("Normalize")
+        dataset = SimCLRDataSetWrapper(batch_size,num_workers,None,"(12, 250)",None,target_folder,[target_folder],None,None,
+                                       mode="linear_evaluation", transformations=transformations, percentage=percentage, folds=folds, t_params=t_params, test=test, ptb_xl_label="label_all")
+    else:
+        if normalize:
+            # always use PTB-XL stats
+            transformations = ["Normalize"]
+            dataset = SimCLRDataSetWrapper(batch_size,num_workers,None,"(12, 250)",None,target_folder,[target_folder],None,None,
+            mode="linear_evaluation", percentage=percentage, folds=folds, test=test, transformations=transformations, ptb_xl_label="label_all")
+        else:
+            dataset = SimCLRDataSetWrapper(batch_size,num_workers,None,"(12, 250)",None,target_folder,[target_folder],None,None,
+                                           mode="linear_evaluation", percentage=percentage, folds=folds, test=test, ptb_xl_label="label_all")
+
+    train_loader, valid_loader = dataset.get_data_loaders()
+    return dataset, train_loader, valid_loader
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    dataset, train_loader, _ = get_dataset(
+        args.batch_size, args.num_workers, args.dataset, folds=args.folds, test=args.test, normalize=args.normalize)
+    _, _, valid_loader = get_dataset(
+        args.batch_size, args.num_workers, args.dataset, folds=args.folds, test=False, normalize=args.normalize)
+    val_idmap = dataset.val_ds_idmap
+    dataset, _, test_loader = get_dataset(
+        args.batch_size, args.num_workers, args.dataset, test=True, normalize=args.normalize)
+    test_idmap = dataset.val_ds_idmap
+    lbl_itos = dataset.lbl_itos
+    tag = "f=" + str(args.folds) + "_" + args.tag
+    tag = tag if args.use_pretrained else "ran_" + tag
+    tag = "eval_" + tag if args.eval_only else tag
+    model_tag = "finetuned" if args.load_finetuned else "ckpt"
+    if args.test_noised:
+        t_params_by_level = {
+            1: {"bw_cmax": 0.05, "em_cmax": 0.25, "pl_cmax": 0.1, "bs_cmax": 0.5},
+            2: {"bw_cmax": 0.1, "em_cmax": 0.5, "pl_cmax": 0.2, "bs_cmax": 1},
+            3: {"bw_cmax": 0.1, "em_cmax": 1, "pl_cmax": 0.2, "bs_cmax": 2},
+            4: {"bw_cmax": 0.2, "em_cmax": 1, "pl_cmax": 0.4, "bs_cmax": 2},
+            5: {"bw_cmax": 0.2, "em_cmax": 1.5, "pl_cmax": 0.4, "bs_cmax": 2.5},
+            6: {"bw_cmax": 0.3, "em_cmax": 2, "pl_cmax": 0.5, "bs_cmax": 3},
+        }
+        if args.noise_level not in t_params_by_level.keys():
+            raise("noise level does not exist")
+        t_params = t_params_by_level[args.noise_level]
+        dataset, _, noise_valid_loader = get_dataset(
+            args.batch_size, args.num_workers, args.dataset, apply_noise=True, t_params=t_params, test=args.test)
+    else:
+        noise_valid_loader = None
+    losses, macros, predss, result_macros, result_macros_agg, test_macros, test_macros_agg, noised_macros, noised_macros_agg = [
+    ], [], [], [], [], [], [], [], []
+    ckpt_epoch_lin=0
+    ckpt_epoch_fin=0
+    if args.f_epochs == 0:
+        save_model_at = os.path.join(os.path.dirname(
+            args.model_file), "n=" + str(args.noise_level) + "_"+tag + "lin_finetuned")
+        filename = os.path.join(os.path.dirname(
+            args.model_file), "n=" + str(args.noise_level) + "_"+tag + "res_lin.pkl")
+    else:
+        save_model_at = os.path.join(os.path.dirname(
+            args.model_file), "n=" + str(args.noise_level) + "_"+tag + "fin_finetuned")
+        filename = os.path.join(os.path.dirname(
+            args.model_file), "n=" + str(args.noise_level) + "_"+tag + "res_fin.pkl")
+
+    model = load_model(
+        args.linear_evaluation, 71, args.use_pretrained or args.load_finetuned, hidden=args.hidden,
+        location=args.model_file, discriminative_lr=args.discriminative_lr, method=args.method)
+    loss_fn, optimizer = configure_optimizer(
+        model, args.batch_size, head_only=True, discriminative_lr=args.discriminative_lr, discriminative_lr_factor=0.1 if args.use_pretrained and args.discriminative_lr else 1)
+    if not args.eval_only:
+        print("train model...")
+        if not isdir(save_model_at):
+            os.mkdir(save_model_at)
+
+        l1, m1, bm, bm_agg, tm, tm_agg, ckpt_epoch_lin, preds = train_model(model, train_loader, valid_loader, test_loader, args.l_epochs, loss_fn,
+                                                                            optimizer, head_only=True, linear_evaluation=args.linear_evaluation, lr_schedule=args.lr_schedule, save_model_at=join(save_model_at, "finetuned.pt"),
+                                                                            val_idmap=val_idmap, test_idmap=test_idmap, lbl_itos=lbl_itos, cpc=(args.method == "cpc"))
+        if bm != 0:
+            print("best macro after head-only training:", bm_agg)
+        l2 = []
+        m2 = []
+        if args.f_epochs != 0:
+            if args.l_epochs != 0:
+                model = load_model(
+                    False, 71, True, hidden=args.hidden,
+                    location=join(save_model_at, "finetuned.pt"), discriminative_lr=args.discriminative_lr, method=args.method)
+            loss_fn, optimizer = configure_optimizer(
+                model, args.batch_size, head_only=False, discriminative_lr=args.discriminative_lr, discriminative_lr_factor=0.1 if args.use_pretrained and args.discriminative_lr else 1)
+            l2, m2, bm, bm_agg, tm, tm_agg, ckpt_epoch_fin, preds = train_model(model, train_loader, valid_loader, test_loader, args.f_epochs, loss_fn,
+                                                                                optimizer, head_only=False, linear_evaluation=False, lr_schedule=args.lr_schedule, save_model_at=join(save_model_at, "finetuned.pt"),
+                                                                                val_idmap=val_idmap, test_idmap=test_idmap, lbl_itos=lbl_itos, cpc=(args.method == "cpc"))
+        losses.append(l1+l2)
+        macros.append(m1+m2)
+        test_macros.append(tm)
+        test_macros_agg.append(tm_agg)
+        result_macros.append(bm)
+        result_macros_agg.append(bm_agg)
+
+    else:
+        preds, eval_macro, eval_macro_agg = evaluate(
+            model, test_loader, test_idmap, lbl_itos, cpc=(args.method == "cpc"))
+        result_macros.append(eval_macro)
+        result_macros_agg.append(eval_macro_agg)
+        if args.verbose:
+            print("macro:", eval_macro)
+    predss.append(preds)
+
+    if noise_valid_loader is not None:
+        _, noise_macro, noise_macro_agg = evaluate(
+            model, noise_valid_loader, val_idmap, lbl_itos)
+        noised_macros.append(noise_macro)
+        noised_macros_agg.append(noise_macro_agg)
+    res = {"filename": filename, "epochs": args.l_epochs+args.f_epochs, "model_location": args.model_location,
+           "losses": losses, "macros": macros, "predss": predss, "result_macros": result_macros, "result_macros_agg": result_macros_agg,
+           "test_macros": test_macros, "test_macros_agg": test_macros_agg, "noised_macros": noised_macros, "noised_macros_agg": noised_macros_agg, "ckpt_epoch_lin": ckpt_epoch_lin, "ckpt_epoch_fin": ckpt_epoch_fin,
+           "discriminative_lr": args.discriminative_lr, "hidden": args.hidden, "lr_schedule": args.lr_schedule,
+           "use_pretrained": args.use_pretrained, "linear_evaluation": args.linear_evaluation, "loaded_finetuned": args.load_finetuned,
+           "eval_only": args.eval_only, "noise_level": args.noise_level, "test_noised": args.test_noised, "normalized": args.normalize}
+    pickle.dump(res, open(filename, "wb"))
+    print("dumped results to", filename)
+    print(res)
+    print("Done!")
+