def get_train_steps(dataset, train_steps, train_epochs, train_batch_size):
  """Determine the number of training steps."""
  num_train_examples = dataset.num_train_examples
  return train_steps or (
      num_train_examples * train_epochs // train_batch_size + 1)