Drop remainder from datasets in text_classifier. This helps deal with…

… issues on TPU training that results in NaN loss.

PiperOrigin-RevId: 625114628

mediapipe/model_maker/python/core/tasks/classifier.py

@@ -81,6 +81,7 @@ def _train_model(
         is_training=True,
         shuffle=self._shuffle,
         preprocess=preprocessor,
+        drop_remainder=True,
     )
     if self._hparams.repeat and self._hparams.steps_per_epoch is None:
       raise ValueError(
@@ -96,7 +97,9 @@ def _train_model(
     validation_dataset = validation_data.gen_tf_dataset(
         batch_size=self._hparams.batch_size,
         is_training=False,
-        preprocess=preprocessor)
+        preprocess=preprocessor,
+        drop_remainder=True,
+    )
     self._model.compile(
         optimizer=self._optimizer,
         loss=self._loss_function,