diff --git a/mediapipe/model_maker/python/vision/object_detector/object_detector.py b/mediapipe/model_maker/python/vision/object_detector/object_detector.py
index 316df85a9..2d1d92ef3 100644
--- 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()