From 524ac3ca61dc165f23a8d6ce29a9ff36d2fa7e98 Mon Sep 17 00:00:00 2001
From: MediaPipe Team <mediapipe-team@google.com>
Date: Fri, 18 Nov 2022 12:45:56 -0800
Subject: [PATCH] Internal change for Model Maker

PiperOrigin-RevId: 489540387
---
 mediapipe/model_maker/python/core/tasks/classifier.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/mediapipe/model_maker/python/core/tasks/classifier.py b/mediapipe/model_maker/python/core/tasks/classifier.py
index 200726864..f376edffa 100644
--- a/mediapipe/model_maker/python/core/tasks/classifier.py
+++ b/mediapipe/model_maker/python/core/tasks/classifier.py
@@ -91,6 +91,10 @@ class Classifier(custom_model.CustomModel):
     self._history = self._model.fit(
         x=train_dataset,
         epochs=self._hparams.epochs,
+        # `steps_per_epoch` is intentionally set to None in case the dataset
+        # is not repeated. Otherwise, the training process will stop when the
+        # dataset is exhausted even if there are epochs remaining.
+        steps_per_epoch=None,
         validation_data=validation_dataset,
         callbacks=self._callbacks)