def METHOD_NAME(self, state, model, cfg):
    for k in list(state["model"].keys()):
        if (
            k.startswith("shared_decoder") or
            k.startswith("_ema") or
            "decoder" in k
        ):
            logger.info(f"Deleting {k} from checkpoint")
            del state["model"][k]
    model.load_state_dict(state["model"], strict=True)