class_docstr: Training state management class.
  
  This class provides apis for backing up and restoring the training state.
  This allows model and epoch information to be saved periodically and restore
  for fault-tolerance, also known as preemption-recovery purpose.
class_docstr_tok: ["Training","state","management","class",".","This","class","provides","apis","for","backing","up","and","restoring","the","training","state",".","This","allows","model","and","epoch","information","to","be","saved","periodically","and","restore","for","fault-tolerance",",","also","known","as","preemption-recovery","purpose","."]
docstr: Back up the current state of training into a checkpoint file.
  
  Args:
    epoch: The current epoch information to be saved.
docstr_tok: ["Back","up","the","current","state","of","training","into","a","checkpoint","file",".","Args",":","epoch",":","The","current","epoch","information","to","be","saved","."]
code: def back_up(self, epoch):
      """Back up the current state of training into a checkpoint file.
  
      Args:
        epoch: The current epoch information to be saved.
      """
      backend.set_value(self._ckpt_saved_epoch, epoch)
      # Save the model plus CKPT_SAVED_EPOCH variable.
      if self.write_checkpoint_manager.save():
          distributed_file_utils.remove_temp_dirpath(
              self.write_checkpoint_manager.directory,
              self._model.distribute_strategy,
          )
  
file_path: tensorflow-upstream/tensorflow/python/keras/distribute/worker_training_state.py
imports: import os
  from tensorflow.python.checkpoint import checkpoint
  from tensorflow.python.checkpoint import checkpoint_management
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import errors
  from tensorflow.python.keras import backend
  from tensorflow.python.keras.distribute import distributed_file_utils
  from tensorflow.python.keras.utils import mode_keys
  from tensorflow.python.lib.io import file_io
  from tensorflow.python.ops import variables
label_desc: Use image node_id 2 for calling the WorkerTrainingState obj's underlying member method code with example usage: obj.back_up(epoch) without return types