from typing import OrderedDict


def filter_state_dict(state_dict, remove_layers=[]):
    """
    filters the inputted state_dict by removing the layers specified
    in remove_layers
    Arguments:
        state_dict (OrderedDict): state_dict of pytorch model
        remove_layers (list(str)): list of layers to remove 
    """

    state_dict = OrderedDict(
        {key:value for key,value in state_dict.items()
        if key not in remove_layers}
        )
    return state_dict