docstr: Load the contents from a checkpoint and restore the state of the given objects.
  
  Args:
      path: A path to where the file is located
      state: A dictionary of objects whose state will be restored in-place from the checkpoint path.
          This should contain exactly one model, and the model must already be set up by DeepSpeed.
      strict: Whether to enforce that the keys in `state` match the keys in the checkpoint.
  
  Returns:
      Dictionary with the state inside DeepSpeed's engine
  
  Raises:
      ValueError:
          If no state is provided, when no :class:`deepspeed.DeepSpeedEngine` objects were found in the
          state, or when multiple :class:`deepspeed.DeepSpeedEngine` objects were found.
      RuntimeError:
          If DeepSpeed was unable to load the checkpoint due to missing files or because the checkpoint is
          not in the expected DeepSpeed format.
docstr_tok: ["Load","the","contents","from","a","checkpoint","and","restore","the","state","of","the","given","objects",".","Args",":","path",":","A","path","to","where","the","file","is","located","state",":","A","dictionary","of","objects","whose","state","will","be","restored","in-place","from","the","checkpoint","path",".","This","should","contain","exactly","one","model",",","and","the","model","must","already","be","set","up","by","DeepSpeed",".","strict",":","Whether","to","enforce","that","the","keys","in","`","state","`","match","the","keys","in","the","checkpoint",".","Returns",":","Dictionary","with","the","state","inside","DeepSpeed","'s","engine","Raises",":","ValueError",":","If","no","state","is","provided",",","when","no",":","class",":","`","deepspeed.DeepSpeedEngine","`","objects","were","found","in","the","state",",","or","when","multiple",":","class",":","`","deepspeed.DeepSpeedEngine","`","objects","were","found",".","RuntimeError",":","If","DeepSpeed","was","unable","to","load","the","checkpoint","due","to","missing","files","or","because","the","checkpoint","is","not","in","the","expected","DeepSpeed","format","."]
code: def load_checkpoint(
      self,
      path: _PATH,
      state: Optional[
          Union[
              Module,
              Optimizer,
              Dict[str, Union[Module, Optimizer, Any]],
          ]
      ] = None,
      strict: bool = True,
  ) -> Dict[str, Any]:
      """Load the contents from a checkpoint and restore the state of the given objects.
  
      Args:
          path: A path to where the file is located
          state: A dictionary of objects whose state will be restored in-place from the checkpoint path.
              This should contain exactly one model, and the model must already be set up by DeepSpeed.
          strict: Whether to enforce that the keys in `state` match the keys in the checkpoint.
  
      Returns:
          Dictionary with the state inside DeepSpeed's engine
  
      Raises:
          ValueError:
              If no state is provided, when no :class:`deepspeed.DeepSpeedEngine` objects were found in the
              state, or when multiple :class:`deepspeed.DeepSpeedEngine` objects were found.
          RuntimeError:
              If DeepSpeed was unable to load the checkpoint due to missing files or because the checkpoint is
              not in the expected DeepSpeed format.
  
      """
      if (
          isinstance(state, (Module, Optimizer))
          or self.load_full_weights
          and self.zero_stage_3
      ):
          # This code path to enables loading a checkpoint from a non-deepspeed checkpoint or from
          # a consolidated checkpoint
          path = self.broadcast(path)
          return super().load_checkpoint(
              path=path, state=state, strict=strict
          )
  
      if not state:
          raise ValueError(
              f"Got DeepSpeedStrategy.load_checkpoint(..., state={state!r}) but a state with at least "
              f" a model instance to reload is required. Pass it in like so:"
              " DeepSpeedStrategy.load_checkpoint(..., state={'model': model, ...})"
          )
      _validate_checkpoint_directory(path)
  
      engines = _get_deepspeed_engines_from_state(state)
      if len(engines) == 0:
          raise ValueError(
              "Could not find a DeepSpeed model in the provided checkpoint state. Please provide the model as"
              " part of the state like so: `load_checkpoint(..., state={'model': model, ...})`. Make sure"
              " you set up the model (and optimizers if any) through the strategy before loading the checkpoint."
          )
      if len(engines) > 1:
          raise ValueError(
              "Found multiple DeepSpeed engine modules in the given state. Saving and loading checkpoints"
              " with DeepSpeed is currently limited to a single model per checkpoint. To load multiple model"
              " states, call the load method for each model checkpoint separately."
          )
      engine = engines[0]
  
      from deepspeed.runtime import DeepSpeedOptimizer
  
      optimzer_state_requested = any(
          isinstance(item, (Optimizer, DeepSpeedOptimizer))
          for item in state.values()
      )
  
      torch.cuda.empty_cache()
      _, client_state = engine.load_checkpoint(
          path,
          tag="checkpoint",
          load_optimizer_states=optimzer_state_requested,
          load_lr_scheduler_states=False,
          load_module_strict=strict,
      )
  
      if client_state is None:
          raise RuntimeError(
              "DeepSpeed was unable to load the checkpoint. Ensure you passed in a DeepSpeed compatible checkpoint"
              " or a single checkpoint file by setting `DeepSpeedStrategy(..., load_full_weights=True)`."
          )
  
      # `Engine.load_checkpoint` adds useless keys 'optimizer' and 'lr_scheduler' to the client state; remove
      # them to avoid name collision with user state
      keys = set(client_state) & set(state) - {
          "optimizer",
          "lr_scheduler",
      }
      _move_state_into(
          source=client_state, destination=state, keys=keys
      )
      return client_state
  
imports: import argparse
  import json
  import logging
  import os
  import platform
  from contextlib import ExitStack
  from itertools import chain
  from pathlib import Path
  from typing import TYPE_CHECKING, Any, Callable, ContextManager, Dict, List, Mapping, Optional, Tuple, Union
  import torch
  from lightning_utilities.core.imports import RequirementCache
  from torch.nn import Module
  from torch.optim import Optimizer
  from typing_extensions import override
  from lightning.fabric.accelerators import Accelerator, CUDAAccelerator
  from lightning.fabric.plugins.environments.cluster_environment import ClusterEnvironment
  from lightning.fabric.plugins.precision import Precision
  from lightning.fabric.strategies.ddp import DDPStrategy
  from lightning.fabric.strategies.registry import _StrategyRegistry
  from lightning.fabric.strategies.strategy import _Sharded
  from lightning.fabric.utilities.distributed import log
  from lightning.fabric.utilities.load import _move_state_into
  from lightning.fabric.utilities.rank_zero import rank_zero_info, rank_zero_warn
  from lightning.fabric.utilities.seed import reset_seed
  from lightning.fabric.utilities.types import _PATH
label_desc: Use image node_id 11 for calling the DeepSpeedStrategy obj's underlying member method code with example usage: obj.load_checkpoint(path, state, strict) and returns: client_state, super