message: Use int64 for `CKPT_SAVED_EPOCH` variable as that's what 'VarHandleOp' OpKernel for GPU devices is registered with.
diff: @@ -23,6 +23,8 @@ import shutil
  import tempfile
  from tensorflow.python.distribute import distribute_coordinator_context as dc_context
  +from tensorflow.python.framework import constant_op
  +from tensorflow.python.framework import dtypes
  from tensorflow.python.keras import backend as K
  from tensorflow.python.keras.utils import mode_keys
  from tensorflow.python.ops import variables
  @@ -78,8 +80,11 @@ class MultiWorkerTrainingState(object):
  self._temp_dir, self._temp_filepath = self._get_temp_filepath()
  # The epoch at which the checkpoint is saved. Used for fault-tolerance.
  + # GPU device only has int64 dtype registered VarHandleOp.
  self._ckpt_saved_epoch = variables.Variable(
  - initial_value=CKPT_SAVED_EPOCH_UNUSED_VALUE, name='ckpt_saved_epoch')
  + initial_value=constant_op.constant(
  + CKPT_SAVED_EPOCH_UNUSED_VALUE, dtype=dtypes.int64),
  + name='ckpt_saved_epoch')
  # Variable initialization.
  K.set_value(self._ckpt_saved_epoch, CKPT_SAVED_EPOCH_UNUSED_VALUE)
  