:zap: Made some more modifications

app.py

@@ -168,7 +168,7 @@ class TorchMetrics:
     self.binary_recall.update(input=batch_predicted_labels, target=batch_actual_labels)
     pass
 
-  def compute_and_reset_on_epoch_end(self, log, log_prefix: str, log_color: str = green):
+  def compute_metrics_and_log(self, log, log_prefix: str, log_color: str = green):
     b_accuracy = self.binary_accuracy.compute()
     b_auc = self.binary_auc.compute()
     b_f1_score = self.binary_f1_score.compute()
@@ -181,13 +181,14 @@ class TorchMetrics:
     log(f"{log_prefix}_precision", b_precision)
     log(f"{log_prefix}_recall", b_recall)
 
+    pass
+
+  def reset_on_epoch_end(self):
     self.binary_accuracy.reset()
     self.binary_auc.reset()
     self.binary_f1_score.reset()
     self.binary_precision.reset()
     self.binary_recall.reset()
-    pass
-
 
 
 class MQtlBertClassifierLightningModule(LightningModule):
@@ -241,12 +242,14 @@ class MQtlBertClassifierLightningModule(LightningModule):
     self.log("train_loss", loss)
     # calculate the scores start
     self.train_metrics.update_on_each_step(batch_predicted_labels=preds.squeeze(), batch_actual_labels=y)
+    self.train_metrics.compute_metrics_and_log(log=self.log, log_prefix="train")
     # self.train_metrics.compute_and_log_on_each_step(log=self.log, log_prefix="train")
     # calculate the scores end
     return loss
 
   def on_train_epoch_end(self) -> None:
-    self.train_metrics.compute_and_reset_on_epoch_end(log=self.log, log_prefix="train")
+    self.train_metrics.compute_metrics_and_log(log=self.log, log_prefix="train")
+    self.train_metrics.reset_on_epoch_end()
     pass
 
   def validation_step(self, batch, batch_idx, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
@@ -259,13 +262,15 @@ class MQtlBertClassifierLightningModule(LightningModule):
     self.log("valid_loss", loss)
     # calculate the scores start
     self.validate_metrics.update_on_each_step(batch_predicted_labels=preds.squeeze(), batch_actual_labels=y)
+    self.validate_metrics.compute_metrics_and_log(log=self.log, log_prefix="validate", log_color=blue)
     # self.validate_metrics.compute_and_log_on_each_step(log=self.log, log_prefix="validate", log_color=blue)
 
     # calculate the scores end
     return loss
 
   def on_validation_epoch_end(self) -> None:
-    self.validate_metrics.compute_and_reset_on_epoch_end(log=self.log, log_prefix="validate", log_color=blue)
+    self.validate_metrics.compute_metrics_and_log(log=self.log, log_prefix="validate", log_color=blue)
+    self.validate_metrics.reset_on_epoch_end()
     return None
 
   def test_step(self, batch, batch_idx, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
@@ -276,13 +281,15 @@ class MQtlBertClassifierLightningModule(LightningModule):
     self.log("test_loss", loss)  # do we need this?
     # calculate the scores start
     self.test_metrics.update_on_each_step(batch_predicted_labels=preds.squeeze(), batch_actual_labels=y)
+    self.test_metrics.compute_metrics_and_log(log=self.log, log_prefix="test", log_color=magenta)
     # self.test_metrics.compute_and_log_on_each_step(log=self.log, log_prefix="test", log_color=magenta)
 
     # calculate the scores end
     return loss
 
   def on_test_epoch_end(self) -> None:
-    self.test_metrics.compute_and_reset_on_epoch_end(log=self.log, log_prefix="test", log_color=magenta)
+    self.test_metrics.compute_metrics_and_log(log=self.log, log_prefix="test", log_color=magenta)
+    self.test_metrics.reset_on_epoch_end()
     return None
 
   pass