For data loader, allow drop_remainer to be set.

PiperOrigin-RevId: 351491512

Author: lintian06

tensorflow_examples/lite/model_maker/core/data_util/dataloader.py

@@ -63,7 +63,8 @@ def gen_dataset(self,
                   is_training=False,
                   shuffle=False,
                   input_pipeline_context=None,
-                  preprocess=None):
+                  preprocess=None,
+                  drop_remainder=False):
     """Generate a shared and batched tf.data.Dataset for training/evaluation.
 
     Args:
@@ -76,6 +77,7 @@ def gen_dataset(self,
         among multiple workers when distribution strategy is used.
       preprocess: A function taking three arguments in order, feature, label and
         boolean is_training.
+      drop_remainder: boolean, whether the finaly batch drops remainder.
 
     Returns:
       A TF dataset ready to be consumed by Keras model.
@@ -100,7 +102,7 @@ def gen_dataset(self,
         ds = ds.shuffle(buffer_size=min(self._size, buffer_size))
       ds = ds.repeat()
 
-    ds = ds.batch(batch_size)
+    ds = ds.batch(batch_size, drop_remainder=drop_remainder)
     ds = ds.prefetch(tf.data.experimental.AUTOTUNE)
     # TODO(b/171449557): Consider converting ds to distributed ds here.
     return ds