tiny update (#216)

* tiny update

* Update pyproject.toml

Author: Sllambias

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "yucca"
-version = "2.2.10"
+version = "2.3.0"
 authors = [
   { name="Sebastian Llambias", email="llambias@live.com" },
   { name="Asbjørn Munk", email="9844416+asbjrnmunk@users.noreply.github.com" },

yucca/modules/callbacks/loggers.py

@@ -37,6 +37,7 @@ def __init__(
         self.previous_epoch = 0
         self.hparams: Dict[str, Any] = {}
         self.NAME_HPARAMS_FILE = "hparams.yaml"
+        self.keys_logged_this_epoch = []
 
         if self.disable_logging is False:
             if self.log_file is None:
@@ -103,6 +104,7 @@ def log_metrics(self, metrics, step):
         with open(self.log_file, "a+") as f:
             current_epoch = (step + 1) // self.steps_per_epoch
             if current_epoch != self.previous_epoch:
+                self.keys_logged_this_epoch = ["epoch"]
                 epoch_end_time = time()
                 f.write("\n")
                 f.write("\n")
@@ -114,10 +116,13 @@ def log_metrics(self, metrics, step):
                 self.previous_epoch = current_epoch
                 self.epoch_start_time = epoch_end_time
             for key in metrics:
-                if key == "epoch":
+                if key in self.keys_logged_this_epoch:
                     continue
-                f.write(f"{t} {key+':':20} {metrics[key]} \n")
-                print(f"{t} {key+':':20} {metrics[key]}")
+                else:
+                    f.write(f"{t} {key+':':20} {metrics[key]} \n")
+                    print(f"{t} {key+':':20} {metrics[key]}")
+                    self.keys_logged_this_epoch.append(key)
+
         sys.stdout.flush()
 
     @rank_zero_only