from typing import OrderedDict


def eliminate_except_param(loaded_param, not_eliminate):
    """
    Ex:
        loaded2 = torch.load(os.path.join('model2-baseline', 'checkpoint.pth'))
        loaded2 = eliminate_except_param(loaded2, ['net.1.weight', 'net.1.bias'])
        loaded.update(loaded2)
    :param loaded_param: An OrderedDict.
    :param not_eliminate: A list of key.
    :return adict: OrderedDict with keys in not_eliminate and their values.

    """
    assert type(not_eliminate) == list
    keep_key = []
    keep_value = []
    for key, value in loaded_param.items():
        if key in not_eliminate:
            keep_key.append(key)
            keep_value.append(value)

    adict = OrderedDict()
    for i, j in zip(keep_key, keep_value):
        adict[i] = j
    return adict