def get_checkpoint_steps(dataset, checkpoint_steps,
                         checkpoint_epochs, train_batch_size):
  """Determine the number of checkpoint steps."""
  num_train_examples = dataset.num_train_examples
  return checkpoint_steps or checkpoint_epochs * int(
      round(num_train_examples / train_batch_size))