From 579f3edc3e20ddbe8207ee0c4189a270b2dfedc1 Mon Sep 17 00:00:00 2001
From: Gustaf Rydholm <gustaf.rydholm@gmail.com>
Date: Sat, 2 Sep 2023 01:52:44 +0200
Subject: Refactor lit models

---
 text_recognizer/model/base.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

(limited to 'text_recognizer/model/base.py')

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)
-- 
cgit v1.2.3-70-g09d2