class_docstr: This class mirrors deepspeed functionality. Not all changes are implemented in this version.
  
  See TorchEngineFull for more modifications.
class_docstr_tok: ["This","class","mirrors","deepspeed","functionality",".","Not","all","changes","are","implemented","in","this","version",".","See","TorchEngineFull","for","more","modifications","."]
docstr: Path, identifier and additional client state. This checkpoint can be used to resume training.
  The default behavior is to save this checkpoint relative to the training working directory.
  
  Has to be .pth because safetensors are annoying to dump a bunch of optim states, scales and schedules
docstr_tok: ["Path",",","identifier","and","additional","client","state",".","This","checkpoint","can","be","used","to","resume","training",".","The","default","behavior","is","to","save","this","checkpoint","relative","to","the","training","working","directory",".","Has","to","be",".pth","because","safetensors","are","annoying","to","dump","a","bunch","of","optim","states",",","scales","and","schedules"]
code: def save_training_checkpoint(
      self, identifier="intermediate.pth", directory="", metadata=None
  ):
      """Path, identifier and additional client state. This checkpoint can be used to resume training.
      The default behavior is to save this checkpoint relative to the training working directory.
  
      Has to be .pth because safetensors are annoying to dump a bunch of optim states, scales and schedules
      """
      file = os.path.join(directory, str(identifier))
      if directory != "":
          os.makedirs(directory, exist_ok=True)
  
      save_state = dict()
      save_state["optim"] = self.optimizer.state_dict()
      save_state[
          "model"
      ] = (
          self.model.state_dict()
      )  # this is the raw key containing _orig and _module flags
      save_state["scheduler"] = self.scheduler.state_dict()
      save_state["scaler"] = self.scaler.state_dict()
      save_state["metadata"] = metadata
      torch.save(save_state, file)
  
imports: import torch
  import torch._inductor.utils
  import os
  import json
  from omegaconf import OmegaConf
  from functools import partial
  from contextlib import nullcontext
  import time
  import logging
  import transformers
  from safetensors.torch import load_file, save_file
  from transformers.utils.generic import working_or_temp_dir
  from torch.distributed.optim import ZeroRedundancyOptimizer
  from .utils import group_parameters, prepare_pretraining_dataloader, update_ema, updated_latest_weight_average
  from .optimizers.schedulers import get_schedule_fn
  from .optimizers import Adahessian, AdamWScale, Shampoo, LARS, SAM, ProgressiveBatching, AGD, Sophia
  import warnings
label_desc: Use image node_id 15 for calling the TorchEngineMinimal obj's underlying member method code with example usage: obj.save_training_checkpoint(identifier, directory, metadata) without return types