From 6605f551e77689b5a1e91ac48114c9bad9ae831f Mon Sep 17 00:00:00 2001 From: MediaPipe Team Date: Wed, 5 Apr 2023 14:09:08 -0700 Subject: [PATCH] Object Detector add batch_size and train_data to get_steps_per_epoch. PiperOrigin-RevId: 522149938 --- .../python/vision/object_detector/object_detector.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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()