message: Make sure to batch the dataset correctly
  Summary:
  Batch the dataset and drop the remainder to
  make sure the sizes are preserved.
  Ref T23733
  Test Plan: All tests pass
  Reviewers: #tensorflow, jakeh, davidn, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved
  Maniphest Tasks: T23733
diff: @@ -93,7 +93,17 @@ def _validate_dataset_element_count(ds, count, fn_name):
  def _get_dataset_and_count(x, y, batch_size):
  adapter_cls = data_adapter.select_data_adapter(x, y)
  adapter = adapter_cls(x, y, batch_size=batch_size)
  - return adapter.get_dataset(), adapter.get_size()
  +
  + dataset = adapter.get_dataset()
  + size = adapter.get_size()
  +
  + if adapter.has_partial_batch():
  + dataset = dataset.unbatch()
  + # Remove the partial batch from the dataset.
  + dataset = dataset.batch(batch_size, drop_remainder=True)
  + size -= 1
  +
  + return dataset, size
  class _TensorflowOptimizerWrapper(Optimizer):
  