code: def load_model_checkpoint(
      load_checkpoint_dir: pathlib.Path,
      model: torch.nn.Module,
      optimizer: torch.optim.Optimizer,
  ) -> Tuple[int, torch.nn.Module, torch.optim.Optimizer]:
      """Loads the optimizer state dict and model state dict from the load_checkpoint_dir
      into the passed model and optimizer. Searches for the most recent checkpoint to
      load from
  
      Args:
          load_checkpoint_dir (pathlib.Path):
              The base checkpoint directory to load from
          model (torch.nn.Module):
              The model to load the checkpoint weights into
          optimizer (torch.optim.Optimizer):
              The optimizer to load the checkpoint weigths into
  
      Returns:
          Tuple[int, torch.nn.Module, torch.optim.Optimizer]:
              The checkpoint step, model with state_dict loaded and
              optimizer with state_dict loaded
  
      """
      log_dist(
          f"Loading model and optimizer checkpoint from {load_checkpoint_dir}",
          ranks=[0],
          level=logging.INFO,
      )
      checkpoint_files = list(
          filter(
              lambda path: re.search(
                  r"iter_(?P<iter_no>\d+)\.pt", path.name
              )
              is not None,
              load_checkpoint_dir.glob("*.pt"),
          )
      )
      assert (
          len(checkpoint_files) > 0
      ), "No checkpoints found in directory"
      checkpoint_files = sorted(
          checkpoint_files,
          key=lambda path: int(
              re.search(r"iter_(?P<iter_no>\d+)\.pt", path.name).group(
                  "iter_no"
              )
          ),
      )
      latest_checkpoint_path = checkpoint_files[-1]
      checkpoint_step = int(
          re.search(
              r"iter_(?P<iter_no>\d+)\.pt", latest_checkpoint_path.name
          ).group("iter_no")
      )
  
      state_dict = torch.load(latest_checkpoint_path)
      model.load_state_dict(state_dict["model"], strict=True)
      optimizer.load_state_dict(state_dict["optimizer"])
      log_dist(
          f"Loading model and optimizer checkpoints done. Loaded from {latest_checkpoint_path}",
          ranks=[0],
          level=logging.INFO,
      )
      return checkpoint_step, model, optimizer
  
imports: import os
  import datetime
  import json
  import pathlib
  import re
  import string
  from functools import partial
  from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, TypeVar, Union
  import random
  import datasets
  import fire
  import logging
  import loguru
  import numpy
  import pytz
  import sh
  import torch
  import torch.nn
  import deepspeed
  from torch.utils.data import DataLoader, Dataset
  from torch.utils.tensorboard import SummaryWriter
  from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast
  from transformers.models.roberta import RobertaConfig, RobertaModel
  from transformers.models.roberta.modeling_roberta import RobertaLMHead, RobertaPreTrainedModel
  from deepspeed.accelerator import get_accelerator
label_desc: Use image node_id 9 for calling a global function with example usage: load_model_checkpoint(load_checkpoint_dir, model, optimizer) and returns: checkpoint_step,  model,  optimizer