def unwrap_checkpoint(m: torch.nn.Module):
    """
    unwrap a module and its children from checkpoint_wrapper
    """
    for module in m.modules():
        if hasattr(module, "precheckpoint_forward"):
            module.forward = module.precheckpoint_forward
            del module.precheckpoint_forward
        if hasattr(module, "old_deepcopy_method"):
            module.__deepcopy__ = module.old_deepcopy_method
            del module.old_deepcopy_method
    return m
unwrap a module and its children from checkpoint_wrapper
fairseq/modules/checkpoint_activations.py