):
      saver.save(repr_ds)


class TfRecordRepresentativeDatasetTest(test.TestCase):
  """Test cases for TfRecordRepresentativeDatasetLoader."""

  def test_tf_record_saver_with_generator_dataset(self):
    tf_record_path = self.create_tempfile().full_path
    path_map = {'serving_default': tf_record_path}
    num_samples = 2

    def data_gen():
      for _ in range(num_samples):