code: def load_state(checkpoint_directory, trial_id):
      checkpoint_directory = pathlib.Path(checkpoint_directory)
  
      with checkpoint_directory.joinpath("checkpoint.pt").open(
          "rb"
      ) as f:
          model = torch.load(f)
      with checkpoint_directory.joinpath("state").open("r") as f:
          epochs_completed, ckpt_trial_id = [
              int(field) for field in f.read().split(",")
          ]
  
      # If trial ID does not match our current trial ID, we'll ignore epochs
      # completed and start training from epoch_idx = 0
      if ckpt_trial_id != trial_id:
          epochs_completed = 0
  
      return model, epochs_completed
  
imports: from __future__ import print_function
  import argparse
  import os
  import pathlib
  import filelock
  import torch
  import torch.distributed
  import torch.nn
  import torch.nn.functional
  import torch.optim
  from torch.nn.parallel import DistributedDataParallel
  from torch.optim.lr_scheduler import StepLR
  from torch.utils.data.distributed import DistributedSampler
  from torchvision import datasets, transforms
  import determined
label_desc: Use image node_id 3 for calling a global function with example usage: load_state(checkpoint_directory, trial_id) and returns: model,  epochs_completed