@@ -146,12 +146,14 @@ class Trainer:
         self.enable_qat = False
         self.qat_params = {}
         self._infinite_train_loader = False
+        self._first_backward = True
 
         # METRICS
         self.loss_logging_items_names = None
         self.train_metrics = None
         self.valid_metrics = None
         self.greater_metric_to_watch_is_better = None
+        self.metric_to_watch = None
 
         # SETTING THE PROPERTIES FROM THE CONSTRUCTOR
         self.experiment_name = experiment_name
@@ -426,7 +428,8 @@ class Trainer:
         # ON FIRST BACKWARD, DERRIVE THE LOGGING TITLES.
         if self.loss_logging_items_names is None or self._first_backward:
             self._init_loss_logging_names(loss_logging_items)
-            self._init_monitored_items()
+            if self.metric_to_watch:
+                self._init_monitored_items()
             self._first_backward = False
 
         if len(loss_logging_items) != len(self.loss_logging_items_names):
@@ -1654,6 +1657,8 @@ class Trainer:
         if use_ema_net and self.ema_model is not None:
             self.net = keep_model
 
+        self._first_backward = True
+
         return test_results
 
     def _validate_epoch(self, epoch: int, silent_mode: bool = False) -> tuple:

@@ -18,6 +18,7 @@ from tests.unit_tests import (
     TrainOptimizerParamsOverride,
     CallTrainTwiceTest,
     ResumeTrainingTest,
+    CallTrainAfterTestTest,
 )
 from tests.end_to_end_tests import TestTrainer
 from tests.unit_tests.phase_delegates_test import ContextMethodsTest
@@ -95,6 +96,7 @@ class CoreUnitTestSuiteRunner:
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(TrainOptimizerParamsOverride))
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(PPYoloETests))
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(ResumeTrainingTest))
+        self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(CallTrainAfterTestTest))
 
     def _add_modules_to_end_to_end_tests_suite(self):
         """

@@ -4,6 +4,7 @@ from tests.unit_tests.factories_test import FactoriesTest
 from tests.unit_tests.optimizer_params_override_test import TrainOptimizerParamsOverride
 from tests.unit_tests.resume_training_test import ResumeTrainingTest
 from tests.unit_tests.strictload_enum_test import StrictLoadEnumTest
+from tests.unit_tests.train_after_test_test import CallTrainAfterTestTest
 from tests.unit_tests.zero_weight_decay_on_bias_bn_test import ZeroWdForBnBiasTest
 from tests.unit_tests.save_ckpt_test import SaveCkptListUnitTest
 from tests.unit_tests.all_architectures_test import AllArchitecturesTest
@@ -41,4 +42,5 @@ __all__ = [
     "TrainOptimizerParamsOverride",
     "CallTrainTwiceTest",
     "ResumeTrainingTest",
+    "CallTrainAfterTestTest",
 ]

import unittest
import torch
from super_gradients import Trainer
from super_gradients.training import models
from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader
from super_gradients.training.metrics import Accuracy


class CallTrainAfterTestTest(unittest.TestCase):
    """
    CallTrainTwiceTest

    Purpose is to call train after test and see nothing crashes. Should be ran with available GPUs (when possible)
    so when calling train again we see there's no change in the model's device.
    """

    def test_call_train_after_test(self):
        trainer = Trainer("test_call_train_after_test")
        dataloader = classification_test_dataloader(batch_size=10)

        model = models.get("resnet18", num_classes=5)
        train_params = {
            "max_epochs": 2,
            "lr_updates": [1],
            "lr_decay_factor": 0.1,
            "lr_mode": "step",
            "lr_warmup_epochs": 0,
            "initial_lr": 0.1,
            "loss": torch.nn.CrossEntropyLoss(),
            "optimizer": "SGD",
            "criterion_params": {},
            "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9},
            "train_metrics_list": [Accuracy()],
            "valid_metrics_list": [Accuracy()],
            "metric_to_watch": "Accuracy",
            "greater_metric_to_watch_is_better": True,
        }
        trainer.test(model=model, test_metrics_list=[Accuracy()], test_loader=dataloader)
        trainer.train(model=model, training_params=train_params, train_loader=dataloader, valid_loader=dataloader)

    def test_call_train_after_test_with_loss(self):
        trainer = Trainer("test_call_train_after_test_with_loss")
        dataloader = classification_test_dataloader(batch_size=10)

        model = models.get("resnet18", num_classes=5)
        train_params = {
            "max_epochs": 2,
            "lr_updates": [1],
            "lr_decay_factor": 0.1,
            "lr_mode": "step",
            "lr_warmup_epochs": 0,
            "initial_lr": 0.1,
            "loss": torch.nn.CrossEntropyLoss(),
            "optimizer": "SGD",
            "criterion_params": {},
            "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9},
            "train_metrics_list": [Accuracy()],
            "valid_metrics_list": [Accuracy()],
            "metric_to_watch": "Accuracy",
            "greater_metric_to_watch_is_better": True,
        }
        trainer.test(model=model, test_metrics_list=[Accuracy()], test_loader=dataloader, loss=torch.nn.CrossEntropyLoss())
        trainer.train(model=model, training_params=train_params, train_loader=dataloader, valid_loader=dataloader)


if __name__ == "__main__":
    unittest.main()