from typing import OrderedDict


def uncleanse_state_dict(state_dict):
    """
    This is an mismatch of expected keys, etc. Issua comes up is saved via cpu, but running on gpu, etc.
    Ex. mismatch: keys not matching
    Expecting: {"module.model.0.weight", ...}
    Received: {"model.0.weight", ...}
    """
    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        nk = "module." + k
        new_state_dict[nk] = v
    return new_state_dict