from typing import OrderedDict


def remove_dataparallel_prefix(state_dict): 
    """Removes dataparallel prefix of layer names in a checkpoint state dictionary."""
    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        name = k[7:] if k[:7] == "module." else k
        new_state_dict[name] = v
    return new_state_dict