def get_num_train_images(hparams):
  """Returns the number of training images according to the dataset."""
  num_images_map = {
      'imagenet': 1281167,
      'cifar10': 50000,
  }
  if hparams.input_data.input_fn not in num_images_map:
    raise ValueError(
        f'Unknown dataset size for input_fn {hparams.input_data.input_fn}')

  num_images = num_images_map[hparams.input_data.input_fn]

  if hparams.input_data.max_samples > 0:
    return min(num_images, hparams.input_data.max_samples)
  return num_images