def state_dict_to_cpu(state_dict):
    """Make a copy of the state dict onto the cpu."""
    # .state_dict() references tensors, so we detach and copy to cpu
    return {key: par.detach().cpu() for key, par in state_dict.items()}