Refactor lit models

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)