def METHOD_NAME(
    self,
    dist_model: nn.Module,
    model: nn.Module,
    optimizer: Optional[Optimizer] = None,
    scheduler: Optional["LRScheduler"] = None,
) -> "Checkpoint":
    from ludwig.utils.checkpoint_utils import MultiNodeCheckpoint
    return MultiNodeCheckpoint(self, model, optimizer, scheduler)