Object Detector add batch_size and train_data to get_steps_per_epoch.

PiperOrigin-RevId: 522149938
2023-04-05 14:09:08 -07:00 · 2023-04-05 14:09:08 -07:00 · 6605f551e7
commit 6605f551e7
parent 5615c1e459
1 changed files with 3 additions and 1 deletions
--- a/mediapipe/model_maker/python/vision/object_detector/object_detector.py
+++ b/mediapipe/model_maker/python/vision/object_detector/object_detector.py
@ -105,7 +105,9 @@ class ObjectDetector(classifier.Classifier):
    """
    self._optimizer = self._create_optimizer(
        model_util.get_steps_per_epoch(
-            self._hparams.steps_per_epoch,
+            steps_per_epoch=self._hparams.steps_per_epoch,
+            batch_size=self._hparams.batch_size,
+            train_data=train_data,
        )
    )
    self._create_model()