diff options
Diffstat (limited to 'text_recognizer/model/base.py')
-rw-r--r-- | text_recognizer/model/base.py | 10 |
1 files changed, 10 insertions, 0 deletions
diff --git a/text_recognizer/model/base.py b/text_recognizer/model/base.py index 1cff796..adcb8da 100644 --- a/text_recognizer/model/base.py +++ b/text_recognizer/model/base.py @@ -94,3 +94,13 @@ class LitBase(L.LightningModule): def test_step(self, batch: Tuple[Tensor, Tensor], batch_idx: int) -> None: """Test step.""" pass + + def is_logged_batch(self) -> bool: + if self.trainer is None: + return False + else: + return self.trainer._logger_connector.should_update_logs + + def add_on_first_batch(self, metrics: dict, output: dict, batch_idx: int) -> None: + if batch_idx == 0: + output.update(metrics) |