code: def custom_restore_fn(checkpoint: Checkpoint):
      start = time.perf_counter()
      with checkpoint.as_directory() as checkpoint_dir:
          time_to_load = time.perf_counter() - start
  
          dist.barrier()
          time_tensor = torch.tensor([time_to_load, 1.0])
          dist.reduce(time_tensor, dst=0, op=dist.ReduceOp.SUM)
  
          if train.get_context().get_world_rank() == 0:
              aggregated_metrics = {
                  "load": time_tensor[0].item() / time_tensor[1].item()
              }
              checkpoint.update_metadata(aggregated_metrics)
              print(
                  "[checkpoint] Restore metrics:\n", aggregated_metrics
              )
  
          # This is a file populated by the default saving logic in `train_fn`.
          with open(
              os.path.join(checkpoint_dir, "checkpoint.pkl"), "rb"
          ) as f:
              state = pickle.load(f)
              return state
  
imports: import collections
  from contextlib import contextmanager
  from datetime import datetime
  import json
  import os
  from pathlib import Path
  import pickle
  import shutil
  import subprocess
  import time
  from typing import Any, Dict
  import uuid
  import fsspec
  import numpy
  import pyarrow.fs
  import pytest
  import torch
  import torch.distributed
  import ray
  from ray import train
  from ray._private.dict import flatten_dict
  from ray.air.constants import TRAINING_ITERATION
  from ray.air._internal.uri_utils import URI
  from ray.train import Checkpoint
  from ray.train.base_trainer import TrainingFailedError
  from ray.train.torch import TorchTrainer
  from test_new_persistence import train_fn, _assert_storage_contents, _resume_from_checkpoint
label_desc: Use image node_id 3 for calling a global function with example usage: custom_restore_fn(checkpoint) and returns: state