@@ -124,14 +124,14 @@ The most simple and straightforward way to start training SOTA performance model
 python -m super_gradients.train_from_recipe --config-name=imagenet_regnetY architecture=regnetY800 dataset_interface.data_dir=<YOUR_Imagenet_LOCAL_PATH> ckpt_root_dir=<CHEKPOINT_DIRECTORY>
 ```
 ### Quickly Load Pre-Trained Weights for Your Desired Model with SOTA Performance
-Want to try our pre-trained models on your machine? Import SuperGradients, initialize your SgModel, and load your desired architecture and pre-trained weights from our [SOTA model zoo](#computer-vision-models---pretrained-checkpoints)
-    
+Want to try our pre-trained models on your machine? Import SuperGradients, initialize your Trainer, and load your desired architecture and pre-trained weights from our [SOTA model zoo](#computer-vision-models---pretrained-checkpoints)
+
 ```python
 # The pretrained_weights argument will load a pre-trained architecture on the provided dataset
 # This is an example of loading COCO-2017 pre-trained weights for a YOLOX Nano object detection model
     
 import super_gradients
-from super_gradients.training import SgModel
+from super_gradients.training import Trainer
 
 trainer = SgModel(experiment_name="yoloxn_coco_experiment",ckpt_root_dir=<CHECKPOINT_DIRECTORY>)
 trainer.build_model(architecture="yolox_n", arch_params={"pretrained_weights": "coco", num_classes": 80})

@@ -68,11 +68,11 @@ In this section we present the modifications needed in order to launch your trai
 
 #### Integrating Your Training Code: Main components:
 
-<span style="text-decoration:underline;">SgModel </span>- the main class in charge of training, testing, logging and basically everything that has to do with the execution of training code.
+<span style="text-decoration:underline;">Trainer </span>- the main class in charge of training, testing, logging and basically everything that has to do with the execution of training code.
 
-<span style="text-decoration:underline;">DatasetInterface</span> - which is passed as an argument to the SgModel and wraps the training set, validation set and optionally a test set for the SgModel instance to work with accordingly.
+<span style="text-decoration:underline;">DatasetInterface</span> - which is passed as an argument to the Trainer and wraps the training set, validation set and optionally a test set for the Trainer instance to work with accordingly.
 
-<span style="text-decoration:underline;">SgModel.net</span> -The network to be used for training/testing (of torch.nn.Module type).
+<span style="text-decoration:underline;">Trainer.net</span> -The network to be used for training/testing (of torch.nn.Module type).
 
 
 #### Integrating Your Training Code - Complete Walkthrough: Dataset
@@ -241,13 +241,13 @@ class Top5(torchmetrics.Accuracy):
 
 #### Integrating Your Training Code- Complete Walkthrough: Training script
 
-We instantiate an SgModel and a UserDatasetInterface, then call connect_dataset_interface which will initialize the dataloaders and pass additional dataset parameters to the SgModel instance.
+We instantiate an Trainer and a UserDatasetInterface, then call connect_dataset_interface which will initialize the dataloaders and pass additional dataset parameters to the Trainer instance.
 
 
 ```
-from super_gradients.training import SgModel
+from super_gradients.training import Trainer
 
-sg_model = SgModel(experiment_name='LeNet_cifar10_example')
+sg_model = Trainer(experiment_name='LeNet_cifar10_example')
 dataset_params = {"batch_size": 256}
 dataset = UserDataset(dataset_params)
 sg_model.connect_dataset_interface(dataset)
@@ -255,7 +255,7 @@ sg_model.connect_dataset_interface(dataset)
 ```
 
 
-**Now, we pass a LeNet instance we defined above to the SgModel:**
+**Now, we pass a LeNet instance we defined above to the Trainer:**
 
 
 ```
@@ -589,11 +589,11 @@ You can also connect your training logs into Weights and Biases (WandB) assuming
     **2nd option:**
 
 
-        Set the “launch_tensorboard_process” flag in your training_params passed to SgModel.train(...), and follow instructions displayed in the shell.
+        Set the “launch_tensorboard_process” flag in your training_params passed to Trainer.train(...), and follow instructions displayed in the shell.
 
 
 * **To resume training –**
-When building the network- call SgModel.build_model(...arch_params={'load_checkpoint'True...}). Doing so, will load the network’s weights, as well as any relevant information for resuming training (monitored metric values, optimizer states, etc) with the latest checkpoint. For more advanced usage see SgModel.build_model docs in code.
+When building the network- call Trainer.build_model(...arch_params={'load_checkpoint'True...}). Doing so, will load the network’s weights, as well as any relevant information for resuming training (monitored metric values, optimizer states, etc) with the latest checkpoint. For more advanced usage see Trainer.build_model docs in code.
 
 
 
@@ -613,7 +613,7 @@ When building the network- call SgModel.build_model(...arch_params={'load_checkp
 
 ## Dataset Parameters
 
-dataset_params argument passed to SgModel.build_model().
+dataset_params argument passed to Trainer.build_model().
 
 `batch_size`: int (default=64)
 
@@ -637,13 +637,13 @@ The remote s3 link from which to download the data (optional).
 
 ## Network Architectures
 
-The following architectures are implemented in SuperGradients’ code, and can be initialized by passing their name (i.e string) to SgModel.build_model easily.
+The following architectures are implemented in SuperGradients’ code, and can be initialized by passing their name (i.e string) to Trainer.build_model easily.
 
 For example:
 
 
 ```
-sg_model = SgModel("resnet50_experiment")
+sg_model = Trainer("resnet50_experiment")
 sg_model.build_model(architecture="resnet50")
 ```
 
@@ -894,7 +894,7 @@ Example- how to load a pretrained model:
 
 
 ```
-sg_model = SgModel("resnet50_experiment")
+sg_model = Trainer("resnet50_experiment")
 
 sg_model.build_model(architecture="resnet50",
                       arch_params={"pretrained_weights": "imagenet", "num_classes": 1000}

@@ -1,14 +1,15 @@
 from super_gradients.training import ARCHITECTURES, losses, utils, datasets_utils, DataAugmentation, \
-    TestDatasetInterface, SegmentationTestDatasetInterface, DetectionTestDatasetInterface, ClassificationTestDatasetInterface, SgModel, KDModel
+    TestDatasetInterface, SegmentationTestDatasetInterface, DetectionTestDatasetInterface, ClassificationTestDatasetInterface, SgModel, KDModel, \
+    Trainer, KDTrainer
 from super_gradients.common import init_trainer, is_distributed
 from super_gradients.examples.train_from_recipe_example import train_from_recipe
 from super_gradients.examples.train_from_kd_recipe_example import train_from_kd_recipe
 from super_gradients.sanity_check import env_sanity_check
 
 __all__ = ['ARCHITECTURES', 'losses', 'utils', 'datasets_utils', 'DataAugmentation',
-           'TestDatasetInterface', 'SgModel', 'KDModel', 'SegmentationTestDatasetInterface', 'DetectionTestDatasetInterface',
+           'TestDatasetInterface', 'Trainer', 'KDTrainer', 'SegmentationTestDatasetInterface', 'DetectionTestDatasetInterface',
            'ClassificationTestDatasetInterface', 'init_trainer', 'is_distributed', 'train_from_recipe', 'train_from_kd_recipe',
-           'env_sanity_check']
+           'env_sanity_check', 'KDModel', 'SgModel']
 
 
 env_sanity_check()

@@ -5,7 +5,7 @@ class EvaluationType(str, Enum):
     """
     EvaluationType
 
-    Passed to SgModel.evaluate(..), and controls which phase callbacks should be triggered (if at all).
+    Passed to Trainer.evaluate(..), and controls which phase callbacks should be triggered (if at all).
 
         Attributes:
             TEST

@@ -15,7 +15,7 @@ from super_gradients.common import ADNNModelRepositoryDataInterfaces
 from super_gradients.common.abstractions.abstract_logger import get_logger
 from super_gradients.common.sg_loggers.abstract_sg_logger import AbstractSGLogger
 from super_gradients.common.environment.env_helpers import multi_process_safe
-from super_gradients.training.utils import sg_model_utils
+from super_gradients.training.utils import sg_trainer_utils
 from super_gradients.training.params import TrainingParams
 
 logger = get_logger(__name__)
@@ -89,11 +89,11 @@ class BaseSGLogger(AbstractSGLogger):
 
     @multi_process_safe
     def _launch_tensorboard(self, port):
-        self.tensor_board_process, _ = sg_model_utils.launch_tensorboard_process(self._local_dir, port=port)
+        self.tensor_board_process, _ = sg_trainer_utils.launch_tensorboard_process(self._local_dir, port=port)
 
     @multi_process_safe
     def _init_tensorboard(self, resumed, tb_files_user_prompt):
-        self.tensorboard_writer = sg_model_utils.init_summary_writer(self._local_dir, resumed, tb_files_user_prompt)
+        self.tensorboard_writer = sg_trainer_utils.init_summary_writer(self._local_dir, resumed, tb_files_user_prompt)
 
     @multi_process_safe
     def _make_dir(self):

@@ -11,7 +11,7 @@ Cifar10 training with SuperGradients training with the following initialized tor
 Main purpose is to demonstrate training in SG with minimal abstraction and maximal flexibility
 """
 
-from super_gradients import SgModel
+from super_gradients import Trainer
 from super_gradients.training.metrics.classification_metrics import Accuracy, Top5
 from super_gradients.training import MultiGPUMode
 from torch.optim import ASGD
@@ -48,10 +48,10 @@ loss_fn = CrossEntropyLoss()
 phase_callbacks = [LRSchedulerCallback(scheduler=rop_lr_scheduler, phase=Phase.VALIDATION_EPOCH_END, metric_name="Accuracy"),
                    LRSchedulerCallback(scheduler=step_lr_scheduler, phase=Phase.TRAIN_EPOCH_END)]
 
-# Bring everything together with SgModel and start training
-model = SgModel("Cifar10_external_objects_example", multi_gpu=MultiGPUMode.OFF,
-                train_loader=train_loader, valid_loader=valid_loader, classes=train_dataset.classes)
-model.build_model(net)
+# Bring everything together with Trainer and start training
+trainer = Trainer("Cifar10_external_objects_example", multi_gpu=MultiGPUMode.OFF,
+                  train_loader=train_loader, valid_loader=valid_loader, classes=train_dataset.classes)
+trainer.build_model(net)
 
 train_params = {"max_epochs": 300,
                 "phase_callbacks": phase_callbacks,
@@ -65,4 +65,4 @@ train_params = {"max_epochs": 300,
                 "greater_metric_to_watch_is_better": True,
                 "lr_scheduler_step_type": "epoch"}
 
-model.train(training_params=train_params)
+trainer.train(training_params=train_params)

@@ -17,13 +17,12 @@ import torch
 from super_gradients.training.datasets.dataset_interfaces.dataset_interface import ImageNetDatasetInterface
 
 import super_gradients
-from super_gradients.training import SgModel, MultiGPUMode
+from super_gradients.training import Trainer, MultiGPUMode
 from super_gradients.training.models import HpmStruct
 import argparse
 
 from super_gradients.training.metrics import Accuracy, Top5
 
-
 parser = argparse.ArgumentParser()
 super_gradients.init_trainer()
 
@@ -60,16 +59,16 @@ dataset_params = {"batch_size": args.batch,
                   "auto_augment_config_string": 'rand-m9-mstd0.5'
                   }
 
-model = SgModel(experiment_name=args.experiment_name,
-                multi_gpu=MultiGPUMode.DISTRIBUTED_DATA_PARALLEL if distributed else MultiGPUMode.DATA_PARALLEL,
-                device='cuda')
+trainer = Trainer(experiment_name=args.experiment_name,
+                  multi_gpu=MultiGPUMode.DISTRIBUTED_DATA_PARALLEL if distributed else MultiGPUMode.DATA_PARALLEL,
+                  device='cuda')
 
 dataset = ImageNetDatasetInterface(dataset_params=dataset_params)
 
-model.connect_dataset_interface(dataset, data_loader_num_workers=8 * devices)
+trainer.connect_dataset_interface(dataset, data_loader_num_workers=8 * devices)
 
 arch_params = HpmStruct(**{"num_classes": 1000, "aux_head": False, "classification_mode": True, 'dropout_prob': 0.3})
 
-model.build_model(architecture="ddrnet_23_slim" if args.slim else "ddrnet_23",
-                  arch_params=arch_params)
-model.train(training_params=train_params_ddr)
+trainer.build_model(architecture="ddrnet_23_slim" if args.slim else "ddrnet_23",
+                    arch_params=arch_params)
+trainer.train(training_params=train_params_ddr)

@@ -4,7 +4,7 @@ Deci-lab model export example.
 The main purpose of this code is to demonstrate how to upload the model to the platform, optimize and download it
  after training is complete, using DeciPlatformCallback.
 """
-from super_gradients import SgModel, ClassificationTestDatasetInterface
+from super_gradients import Trainer, ClassificationTestDatasetInterface
 from super_gradients.training.metrics import Accuracy, Top5
 from super_gradients.training.utils.callbacks import DeciLabUploadCallback, ModelConversionCheckCallback
 from deci_lab_client.models import (
@@ -23,15 +23,15 @@ def main(architecture_name: str):
 
     auth_token = "YOUR_API_TOKEN_HERE"
 
-    model = SgModel(
+    trainer = Trainer(
         f"lab_optimization_{architecture_name}_example",
         model_checkpoints_location="local",
         ckpt_root_dir=checkpoint_dir,
     )
     dataset = ClassificationTestDatasetInterface(dataset_params={"batch_size": 10})
-    model.connect_dataset_interface(dataset, data_loader_num_workers=0)
+    trainer.connect_dataset_interface(dataset, data_loader_num_workers=0)
 
-    model.build_model(architecture=architecture_name, arch_params={"use_aux_heads": True, "aux_head": True})
+    trainer.build_model(architecture=architecture_name, arch_params={"use_aux_heads": True, "aux_head": True})
 
     # CREATE META-DATA, AND OPTIMIZATION REQUEST FORM FOR DECI PLATFORM POST TRAINING CALLBACK
     model_name = f"{architecture_name}_for_deci_lab_export_example"
@@ -91,7 +91,7 @@ def main(architecture_name: str):
 
     # RUN TRAINING. ONCE ALL EPOCHS ARE DONE THE OPTIMIZED MODEL FILE WILL BE LOCATED IN THE EXPERIMENT'S
     # CHECKPOINT DIRECTORY
-    model.train(train_params)
+    trainer.train(train_params)
 
 
 if __name__ == "__main__":

@@ -1,12 +1,12 @@
 import os
-from super_gradients.training import SgModel
+from super_gradients.training import Trainer
 from super_gradients.training.datasets.dataset_interfaces import Cifar10DatasetInterface
 from super_gradients.training.metrics.classification_metrics import Accuracy, Top5
 
 os.environ["DECI_PLATFORM_TOKEN"] = "XXX"  # Replace XXX with your token
 
 
-trainer = SgModel(experiment_name='demo-deci-platform-logger')
+trainer = Trainer(experiment_name='demo-deci-platform-logger')
 dataset = Cifar10DatasetInterface(dataset_params={"batch_size": 256, "val_batch_size": 512})
 trainer.connect_dataset_interface(dataset, data_loader_num_workers=8)
 trainer.build_model("resnet18")

@@ -1,7 +1,7 @@
 # Cifar10 Classification Training:
 # Reaches ~94.9 Accuracy after 250 Epochs
 import super_gradients
-from super_gradients import SgModel
+from super_gradients import Trainer
 from super_gradients.training.datasets.dataset_interfaces.dataset_interface import Cifar10DatasetInterface
 from super_gradients.training.metrics.classification_metrics import Accuracy, Top5
 from super_gradients.training.utils.early_stopping import EarlyStop
@@ -20,12 +20,12 @@ train_params = {"max_epochs": 250, "lr_updates": [100, 150, 200], "lr_decay_fact
                 "greater_metric_to_watch_is_better": True, "phase_callbacks": [early_stop_acc, early_stop_val_loss]}
 
 # Define Model
-model = SgModel("Callback_Example")
+trainer = Trainer("Callback_Example")
 
 # Connect Dataset
 dataset = Cifar10DatasetInterface()
-model.connect_dataset_interface(dataset, data_loader_num_workers=8)
+trainer.connect_dataset_interface(dataset, data_loader_num_workers=8)
 
 # Build Model
-model.build_model("resnet18_cifar")
-model.train(training_params=train_params)
+trainer.build_model("resnet18_cifar")
+trainer.train(training_params=train_params)

@@ -13,13 +13,13 @@ def train(cfg: DictConfig) -> None:
     cfg = hydra.utils.instantiate(cfg)
 
     # CONNECT THE DATASET INTERFACE WITH DECI MODEL
-    cfg.sg_model.connect_dataset_interface(cfg.dataset_interface, data_loader_num_workers=cfg.data_loader_num_workers)
+    cfg.trainer .connect_dataset_interface(cfg.dataset_interface, data_loader_num_workers=cfg.data_loader_num_workers)
 
     # BUILD NETWORK
-    cfg.sg_model.build_model(cfg.architecture, load_checkpoint=cfg.load_checkpoint)
+    cfg.trainer .build_model(cfg.architecture, load_checkpoint=cfg.load_checkpoint)
 
     # TRAIN
-    cfg.sg_model.train(training_params=cfg.training_params)
+    cfg.trainer .train(training_params=cfg.training_params)
 
 
 if __name__ == "__main__":

@@ -1,6 +1,6 @@
 # Darknet53 Backbone Training on HAM10000 Dataset
 from super_gradients.training import MultiGPUMode
-from super_gradients.training import SgModel
+from super_gradients.training import Trainer
 from super_gradients.training.datasets.dataset_interfaces.dataset_interface import ClassificationDatasetInterface
 
 # Define Parameters
@@ -11,10 +11,10 @@ arch_params = {'backbone_mode': False, 'num_classes': 7}
 dataset_params = {"batch_size": 16, "test_batch_size": 16, 'dataset_dir': '/data/HAM10000'}
 
 # Define Model
-model = SgModel("Darknet53_Backbone_HAM10000",
-                model_checkpoints_location='local',
-                device='cuda',
-                multi_gpu=MultiGPUMode.DATA_PARALLEL)
+trainer = Trainer("Darknet53_Backbone_HAM10000",
+                  model_checkpoints_location='local',
+                  device='cuda',
+                  multi_gpu=MultiGPUMode.DATA_PARALLEL)
 
 # Connect Dataset
 dataset = ClassificationDatasetInterface(normalization_mean=(0.7483, 0.5154, 0.5353),
@@ -22,10 +22,10 @@ dataset = ClassificationDatasetInterface(normalization_mean=(0.7483, 0.5154, 0.5
                                          resolution=416,
                                          dataset_params=dataset_params)
 
-model.connect_dataset_interface(dataset, data_loader_num_workers=8)
+trainer.connect_dataset_interface(dataset, data_loader_num_workers=8)
 
 # Build Model
-model.build_model("darknet53", arch_params=arch_params)
+trainer.build_model("darknet53", arch_params=arch_params)
 
 # Start Training
-model.train(training_params=train_params)
+trainer.train(training_params=train_params)

@@ -17,13 +17,13 @@ def train(cfg: DictConfig) -> None:
     cfg = hydra.utils.instantiate(cfg)
 
     # CONNECT THE DATASET INTERFACE WITH DECI MODEL
-    cfg.sg_model.connect_dataset_interface(cfg.dataset_interface, data_loader_num_workers=cfg.data_loader_num_workers)
+    cfg.trainer .connect_dataset_interface(cfg.dataset_interface, data_loader_num_workers=cfg.data_loader_num_workers)
 
     # BUILD NETWORK
-    cfg.sg_model.build_model(cfg.architecture, load_checkpoint=cfg.load_checkpoint)
+    cfg.trainer .build_model(cfg.architecture, load_checkpoint=cfg.load_checkpoint)
 
     # TRAIN
-    cfg.sg_model.train(training_params=cfg.training_params)
+    cfg.trainer .train(training_params=cfg.training_params)
 
 
 if __name__ == "__main__":

@@ -16,13 +16,13 @@ def train(cfg: DictConfig) -> None:
     cfg = hydra.utils.instantiate(cfg)
 
     # CONNECT THE DATASET INTERFACE WITH DECI MODEL
-    cfg.sg_model.connect_dataset_interface(cfg.dataset_interface, data_loader_num_workers=cfg.data_loader_num_workers)
+    cfg.trainer .connect_dataset_interface(cfg.dataset_interface, data_loader_num_workers=cfg.data_loader_num_workers)
 
     # BUILD NETWORK
-    cfg.sg_model.build_model(cfg.architecture, load_checkpoint=cfg.load_checkpoint)
+    cfg.trainer .build_model(cfg.architecture, load_checkpoint=cfg.load_checkpoint)
 
     # TRAIN
-    cfg.sg_model.train(training_params=cfg.training_params)
+    cfg.trainer .train(training_params=cfg.training_params)
 
 
 if __name__ == "__main__":

@@ -15,13 +15,13 @@ def train(cfg: DictConfig) -> None:
     cfg = hydra.utils.instantiate(cfg)
 
     # CONNECT THE DATASET INTERFACE WITH DECI MODEL
-    cfg.sg_model.connect_dataset_interface(cfg.dataset_interface, data_loader_num_workers=cfg.data_loader_num_workers)
+    cfg.trainer .connect_dataset_interface(cfg.dataset_interface, data_loader_num_workers=cfg.data_loader_num_workers)
 
     # BUILD NETWORK
-    cfg.sg_model.build_model(cfg.architecture, arch_params=cfg.arch_params, load_checkpoint=cfg.load_checkpoint)
+    cfg.trainer .build_model(cfg.architecture, arch_params=cfg.arch_params, load_checkpoint=cfg.load_checkpoint)
 
     # TRAIN
-    cfg.sg_model.train(training_params=cfg.training_params)
+    cfg.trainer .train(training_params=cfg.training_params)
 
 
 if __name__ == "__main__":

@@ -10,13 +10,13 @@ def train(cfg: DictConfig) -> None:
     cfg = hydra.utils.instantiate(cfg)
 
     # CONNECT THE DATASET INTERFACE WITH DECI MODEL
-    cfg.sg_model.connect_dataset_interface(cfg.dataset_interface, data_loader_num_workers=cfg.data_loader_num_workers)
+    cfg.trainer .connect_dataset_interface(cfg.dataset_interface, data_loader_num_workers=cfg.data_loader_num_workers)
 
     # BUILD NETWORK
-    cfg.sg_model.build_model(cfg.architecture, arch_params=cfg.arch_params, load_checkpoint=cfg.load_checkpoint)
+    cfg.trainer .build_model(cfg.architecture, arch_params=cfg.arch_params, load_checkpoint=cfg.load_checkpoint)
 
     # TRAIN
-    cfg.sg_model.train(training_params=cfg.training_params)
+    cfg.trainer .train(training_params=cfg.training_params)
 
 
 if __name__ == "__main__":

@@ -28,13 +28,13 @@ def train(cfg: DictConfig) -> None:
     cfg = hydra.utils.instantiate(cfg)
 
     # CONNECT THE DATASET INTERFACE WITH DECI MODEL
-    cfg.sg_model.connect_dataset_interface(cfg.dataset_interface, data_loader_num_workers=cfg.data_loader_num_workers)
+    cfg.trainer .connect_dataset_interface(cfg.dataset_interface, data_loader_num_workers=cfg.data_loader_num_workers)
 
     # BUILD NETWORK
-    cfg.sg_model.build_model(cfg.architecture, arch_params=cfg.arch_params, load_checkpoint=cfg.load_checkpoint)
+    cfg.trainer .build_model(cfg.architecture, arch_params=cfg.arch_params, load_checkpoint=cfg.load_checkpoint)
 
     # TRAIN
-    cfg.sg_model.train(training_params=cfg.training_params)
+    cfg.trainer .train(training_params=cfg.training_params)
 
 
 if __name__ == "__main__":

@@ -26,8 +26,8 @@
 """
 import super_gradients
 import torch.distributed
-from super_gradients.training.sg_model import MultiGPUMode
-from super_gradients.training import SgModel
+from super_gradients.training.sg_trainer import MultiGPUMode
+from super_gradients.training import Trainer
 from super_gradients.training.datasets.dataset_interfaces import ImageNetDatasetInterface
 from super_gradients.common.aws_connection.aws_secrets_manager_connector import AWSSecretsManagerConnector
 from super_gradients.training.metrics.classification_metrics import Accuracy, Top5
@@ -57,12 +57,12 @@ dataset_params = {"batch_size": 128}
 model_repo_bucket_name = AWSSecretsManagerConnector.get_secret_value_for_secret_key(aws_env='research',
                                                                                     secret_name='training_secrets',
                                                                                     secret_key='S3.MODEL_REPOSITORY_BUCKET_NAME')
-model = SgModel("test_checkpoints_resnet_8_gpus",
-                model_checkpoints_location='s3://' + model_repo_bucket_name,
-                multi_gpu=MultiGPUMode.DISTRIBUTED_DATA_PARALLEL
-                )
+trainer = Trainer("test_checkpoints_resnet_8_gpus",
+                  model_checkpoints_location='s3://' + model_repo_bucket_name,
+                  multi_gpu=MultiGPUMode.DISTRIBUTED_DATA_PARALLEL
+                  )
 # FOR AWS
 dataset = ImageNetDatasetInterface(data_dir="/data/Imagenet", dataset_params=dataset_params)
-model.connect_dataset_interface(dataset, data_loader_num_workers=8)
-model.build_model("resnet50")
-model.train(training_params=train_params)
+trainer.connect_dataset_interface(dataset, data_loader_num_workers=8)
+trainer.build_model("resnet50")
+trainer.train(training_params=train_params)

@@ -10,9 +10,9 @@
 # P.S. - Use the relevant training params dict if you are running on TZAG or on V100
 
 import torch
-from super_gradients.training import SgModel, MultiGPUMode
+from super_gradients.training import Trainer, MultiGPUMode
 from super_gradients.training.datasets import CoCoSegmentationDatasetInterface
-from super_gradients.training.sg_model.sg_model import StrictLoad
+from super_gradients.training.sg_trainer.sg_trainer import StrictLoad
 from super_gradients.training.metrics.segmentation_metrics import PixelAccuracy, IoU
 
 model_size_str = '34'
@@ -66,16 +66,16 @@ experiment_name_dataset_suffix = '_coco_seg_' + str(
 
 experiment_name = experiment_name_prefix + model_size_str + experiment_name_dataset_suffix
 
-model = SgModel(experiment_name,
-                multi_gpu=MultiGPUMode.DISTRIBUTED_DATA_PARALLEL,
-                ckpt_name='ckpt_best.pth')
+trainer = Trainer(experiment_name,
+                  multi_gpu=MultiGPUMode.DISTRIBUTED_DATA_PARALLEL,
+                  ckpt_name='ckpt_best.pth')
 
 coco_seg_datasaet_interface = CoCoSegmentationDatasetInterface(dataset_params=coco_seg_dataset_tzag_params,
                                                                cache_labels=False,
                                                                dataset_classes_inclusion_tuples_list=coco_sub_classes_inclusion_tuples_list)
 
-model.connect_dataset_interface(coco_seg_datasaet_interface, data_loader_num_workers=data_loader_num_workers)
-model.build_model('shelfnet' + model_size_str, arch_params=shelfnet_lw_arch_params)
+trainer.connect_dataset_interface(coco_seg_datasaet_interface, data_loader_num_workers=data_loader_num_workers)
+trainer.build_model('shelfnet' + model_size_str, arch_params=shelfnet_lw_arch_params)
 
 print('Training ShelfNet-LW model: ' + experiment_name)
-model.train(training_params=shelfnet_coco_training_params)
+trainer.train(training_params=shelfnet_coco_training_params)

@@ -1,5 +1,5 @@
 from super_gradients.training.datasets.dataset_interfaces.dataset_interface import SuperviselyPersonsDatasetInterface
-from super_gradients.training.sg_model import SgModel
+from super_gradients.training.sg_trainer import Trainer
 from super_gradients.training.metrics import BinaryIOU
 from super_gradients.training.transforms.transforms import ResizeSeg, RandomFlip, RandomRescale, CropImageAndMask, \
     PadShortToCropSize, ColorJitterSeg
@@ -19,15 +19,15 @@ dataset_params = {
 
 dataset_interface = SuperviselyPersonsDatasetInterface(dataset_params)
 
-model = SgModel("regseg48_transfer_learning_old_dice_diff_lrs_head_fixed_50_epochs")
+trainer = Trainer("regseg48_transfer_learning_old_dice_diff_lrs_head_fixed_50_epochs")
 
 # CONNECTING THE DATASET INTERFACE WILL SET SGMODEL'S CLASSES ATTRIBUTE ACCORDING TO SUPERVISELY
-model.connect_dataset_interface(dataset_interface)
+trainer.connect_dataset_interface(dataset_interface)
 
 # THIS IS WHERE THE MAGIC HAPPENS- SINCE SGMODEL'S CLASSES ATTRIBUTE WAS SET TO BE DIFFERENT FROM CITYSCAPES'S, AFTER
 # LOADING THE PRETRAINED REGSET, IT WILL CALL IT'S REPLACE_HEAD METHOD AND CHANGE IT'S SEGMENTATION HEAD LAYER ACCORDING
 # TO OUR BINARY SEGMENTATION DATASET
-model.build_model("regseg48", arch_params={"pretrained_weights": "cityscapes"})
+trainer.build_model("regseg48", arch_params={"pretrained_weights": "cityscapes"})
 
 # DEFINE TRAINING PARAMS. SEE DOCS FOR THE FULL LIST.
 train_params = {"max_epochs": 50,
@@ -55,4 +55,4 @@ train_params = {"max_epochs": 50,
                                                                             last_img_idx_in_batch=4)],
                 }
 
-model.train(train_params)
+trainer.train(train_params)

@@ -18,7 +18,7 @@ Finally, once training is over- we trigger a pos-training callback that will exp
 """
 
 from super_gradients.training.datasets.dataset_interfaces.dataset_interface import ImageNetDatasetInterface
-from super_gradients.training import SgModel, MultiGPUMode
+from super_gradients.training import Trainer, MultiGPUMode
 from super_gradients.training.metrics.classification_metrics import Accuracy
 
 import super_gradients
@@ -27,12 +27,12 @@ from super_gradients.training.utils.quantization_utils import PostQATConversionC
 super_gradients.init_trainer()
 
 dataset = ImageNetDatasetInterface(data_dir="/data/Imagenet", dataset_params={"batch_size": 128})
-model = SgModel("resnet18_qat_example",
-                model_checkpoints_location='local',
-                multi_gpu=MultiGPUMode.DISTRIBUTED_DATA_PARALLEL)
+trainer = Trainer("resnet18_qat_example",
+                  model_checkpoints_location='local',
+                  multi_gpu=MultiGPUMode.DISTRIBUTED_DATA_PARALLEL)
 
-model.connect_dataset_interface(dataset)
-model.build_model("resnet18", checkpoint_params={"pretrained_weights": "imagenet"})
+trainer.connect_dataset_interface(dataset)
+trainer.build_model("resnet18", checkpoint_params={"pretrained_weights": "imagenet"})
 
 train_params = {"max_epochs": 1,
                 "lr_mode": "step",
@@ -53,9 +53,9 @@ train_params = {"max_epochs": 1,
                     # statistics method for amax computation (one of [percentile, mse, entropy, max]).
                     "calibrate": True,  # whether to perform calibration.
                     "num_calib_batches": 2,  # number of batches to collect the statistics from.
-                    "percentile": 99.99  # percentile value to use when SgModel,
+                    "percentile": 99.99  # percentile value to use when Trainer,
                 },
                 "phase_callbacks": [PostQATConversionCallback(dummy_input_size=(1, 3, 224, 224))]
                 }
 
-model.train(training_params=train_params)
+trainer.train(training_params=train_params)

@@ -2,8 +2,8 @@
 import super_gradients
 import torch
 from super_gradients.training.datasets import PascalAUG2012SegmentationDataSetInterface
-from super_gradients.training import SgModel, MultiGPUMode
-from super_gradients.training.sg_model.sg_model import StrictLoad
+from super_gradients.training import Trainer, MultiGPUMode
+from super_gradients.training.sg_trainer.sg_trainer import StrictLoad
 from super_gradients.training.metrics.segmentation_metrics import PixelAccuracy, IoU
 
 super_gradients.init_trainer()
@@ -44,13 +44,13 @@ model_size_str = '34'
 experiment_name_prefix = 'shelfnet_lw_'
 experiment_name_dataset_suffix = '_pascal_aug_encoding_dataset_train_250_epochs_no_batchnorm_decoder'
 experiment_name = experiment_name_prefix + model_size_str + experiment_name_dataset_suffix
-model = SgModel(experiment_name, model_checkpoints_location='local', multi_gpu=True,
-                ckpt_name='resnet' + model_size_str + '.pth')
+trainer = Trainer(experiment_name, model_checkpoints_location='local', multi_gpu=True,
+                  ckpt_name='resnet' + model_size_str + '.pth')
 
 pascal_aug_datasaet_interface = PascalAUG2012SegmentationDataSetInterface(
     dataset_params=pascal_aug_dataset_params,
     cache_labels=False)
-model.connect_dataset_interface(pascal_aug_datasaet_interface, data_loader_num_workers=data_loader_num_workers)
-model.build_model('shelfnet' + model_size_str, arch_params=shelfnet_lw_arch_params, checkpoint_params=checkpoint_params)
+trainer.connect_dataset_interface(pascal_aug_datasaet_interface, data_loader_num_workers=data_loader_num_workers)
+trainer.build_model('shelfnet' + model_size_str, arch_params=shelfnet_lw_arch_params, checkpoint_params=checkpoint_params)
 print('Training ShelfNet-LW model: ' + experiment_name)
-model.train(training_params=shelfnet_lw_pascal_aug_training_params)
+trainer.train(training_params=shelfnet_lw_pascal_aug_training_params)

@@ -14,7 +14,7 @@ from super_gradients.training.kd_trainer import KDTrainer
 
 @hydra.main(config_path=pkg_resources.resource_filename("super_gradients.recipes", ""))
 def main(cfg: DictConfig) -> None:
-    KDTrainer.train(cfg)
+    KDTrainer.train_from_config(cfg)
 
 
 if __name__ == "__main__":

@@ -9,12 +9,13 @@ import super_gradients
 from omegaconf import DictConfig
 import hydra
 import pkg_resources
-from super_gradients.training.trainer import Trainer
+
+from super_gradients import Trainer
 
 
 @hydra.main(config_path=pkg_resources.resource_filename("super_gradients.recipes", ""))
 def main(cfg: DictConfig) -> None:
-    Trainer.train(cfg)
+    Trainer.train_from_config(cfg)
 
 
 if __name__ == "__main__":

@@ -1,4 +1,4 @@
-from super_gradients.training import SgModel
+from super_gradients.training import Trainer
 from super_gradients.training import MultiGPUMode
 from dataset import UserDataset
 from model import ResNet, BasicBlock
@@ -11,17 +11,17 @@ def main():
     arch_params = {'num_classes': 10}
     model = ResNet(BasicBlock, [2, 2, 2, 2], num_classes=arch_params['num_classes'])
 
-    deci_classification_model = SgModel('client_model_training',
-                                        model_checkpoints_location='local',
-                                        multi_gpu=MultiGPUMode.OFF)
+    trainer = Trainer('client_model_training',
+                      model_checkpoints_location='local',
+                      multi_gpu=MultiGPUMode.OFF)
 
     # if a torch.nn.Module is provided when building the model, the model will be integrated into deci model class
-    deci_classification_model.build_model(model, arch_params=arch_params)
+    trainer.build_model(model, arch_params=arch_params)
 
     # ------------------ Loading The Dataset From Dataset.py----------------
     dataset_params = {"batch_size": 256}
     dataset = UserDataset(dataset_params)
-    deci_classification_model.connect_dataset_interface(dataset)
+    trainer.connect_dataset_interface(dataset)
 
     # ------------------ Loading The Loss From Loss.py -----------------
     loss = LabelSmoothingCrossEntropyLoss()
@@ -48,7 +48,7 @@ def main():
                     "metric_to_watch": "Accuracy",
                     "greater_metric_to_watch_is_better": True}
 
-    deci_classification_model.train(train_params)
+    trainer.train(train_params)
 
 
 if __name__ == '__main__':

@@ -3,5 +3,5 @@ load_backbone: False # whether to load only backbone part of checkpoint
 external_checkpoint_path: # checkpoint path that is not located in super_gradients/checkpoints
 source_ckpt_folder_name: # dirname for checkpoint loading
 strict_load: # key matching strictness for loading checkpoint's weights
-  _target_: super_gradients.training.sg_model.StrictLoad
+  _target_: super_gradients.training.sg_trainer.StrictLoad
   value: True

@@ -25,11 +25,5 @@ checkpoint_params:
 model_checkpoints_location: local
 ckpt_root_dir:
 
-sg_model:
-  _target_: super_gradients.SgModel
-  experiment_name: ${experiment_name}
-  model_checkpoints_location: ${model_checkpoints_location}
-  multi_gpu: Off
-
 architecture: resnet18_cifar
 

@@ -115,10 +115,4 @@ multi_gpu:
   _target_: super_gradients.training.sg_model.MultiGPUMode
   value: 'DDP'
 
-sg_model:
-  _target_: super_gradients.SgModel
-  experiment_name: ${experiment_name}
-  model_checkpoints_location: ${model_checkpoints_location}
-  ckpt_root_dir: ${ckpt_root_dir}
-  multi_gpu: ${multi_gpu}
 

@@ -126,9 +126,3 @@ checkpoint_params:
 
 model_checkpoints_location: local
 ckpt_root_dir:
-
-sg_model:
-  _target_: super_gradients.SgModel
-  experiment_name: ${experiment_name}
-  model_checkpoints_location: ${model_checkpoints_location}
-  multi_gpu: DDP

@@ -64,10 +64,3 @@ multi_gpu:
   _target_: super_gradients.training.sg_model.MultiGPUMode
   value: 'DDP'
 
-sg_model:
-  _target_: super_gradients.SgModel
-  experiment_name: ${experiment_name}
-  model_checkpoints_location: ${model_checkpoints_location}
-  ckpt_root_dir: ${ckpt_root_dir}
-  multi_gpu: ${multi_gpu}
-

@@ -38,11 +38,6 @@ experiment_suffix: res${dataset_params.train_image_size}
 experiment_name: ${architecture}_coco2017_${experiment_suffix}
 
 ckpt_root_dir:
-sg_model:
-  _target_: super_gradients.SgModel
-  experiment_name: ${experiment_name}
-  model_checkpoints_location: ${model_checkpoints_location}
-  multi_gpu: ${multi_gpu}
-  ckpt_root_dir: ${ckpt_root_dir}
+
 
 

@@ -32,17 +32,9 @@ checkpoint_params:
 experiment_name: coco_segmentation_21_subclass_shelfnet34
 
 multi_gpu:
-  _target_: super_gradients.training.sg_model.MultiGPUMode
+  _target_: super_gradients.training.sg_trainer.MultiGPUMode
   value: 'DDP'
 
 ckpt_root_dir:
 
-sg_model:
-  _target_: super_gradients.SgModel
-  experiment_name: ${experiment_name}
-  model_checkpoints_location: ${model_checkpoints_location}
-  ckpt_root_dir: ${ckpt_root_dir}
-  multi_gpu: ${multi_gpu}
-
-
 architecture: shelfnet34_lw

@@ -41,14 +41,7 @@ model_checkpoints_location: local
 ckpt_root_dir:
 
 multi_gpu:
-  _target_: super_gradients.training.sg_model.MultiGPUMode
+  _target_: super_gradients.training.sg_trainer.MultiGPUMode
   value: 'DDP'
 
-sg_model:
-  _target_: super_gradients.SgModel
-  experiment_name: ${experiment_name}
-  model_checkpoints_location: ${model_checkpoints_location}
-  ckpt_root_dir: ${ckpt_root_dir}
-  multi_gpu: ${multi_gpu}
-
 architecture: efficientnet_b0

@@ -42,14 +42,7 @@ experiment_name: mobileNetv2_training
 ckpt_root_dir:
 
 multi_gpu:
-  _target_: super_gradients.training.sg_model.MultiGPUMode
+  _target_: super_gradients.training.sg_trainer.MultiGPUMode
   value: 'DDP'
 
-sg_model:
-  _target_: super_gradients.SgModel
-  experiment_name: ${experiment_name}
-  model_checkpoints_location: ${model_checkpoints_location}
-  ckpt_root_dir: ${ckpt_root_dir}
-  multi_gpu: ${multi_gpu}
-
 architecture: mobilenet_v2

@@ -24,14 +24,7 @@ experiment_name: mobileNetv3_large_training
 ckpt_root_dir:
 
 multi_gpu:
-  _target_: super_gradients.training.sg_model.MultiGPUMode
+  _target_: super_gradients.training.sg_trainer.MultiGPUMode
   value: 'DDP'
 
-sg_model:
-  _target_: super_gradients.SgModel
-  experiment_name: ${experiment_name}
-  model_checkpoints_location: ${model_checkpoints_location}
-  ckpt_root_dir: ${ckpt_root_dir}
-  multi_gpu: ${multi_gpu}
-
 architecture: mobilenet_v3_large

@@ -58,14 +58,8 @@ experiment_name: regnetY800_imagenet
 ckpt_root_dir:
 
 multi_gpu:
-  _target_: super_gradients.training.sg_model.MultiGPUMode
+  _target_: super_gradients.training.sg_trainer.MultiGPUMode
   value: 'Off'
 
-sg_model:
-  _target_: super_gradients.SgModel
-  experiment_name: ${experiment_name}
-  model_checkpoints_location: ${model_checkpoints_location}
-  ckpt_root_dir: ${ckpt_root_dir}
-  multi_gpu: ${multi_gpu}
 
 architecture: regnetY800

@@ -36,14 +36,7 @@ checkpoint_params:
 experiment_name: repvgg_a0_imagenet_reproduce_fix
 
 multi_gpu:
-  _target_: super_gradients.training.sg_model.MultiGPUMode
+  _target_: super_gradients.training.sg_trainer.MultiGPUMode
   value: 'DDP'
 
-sg_model:
-  _target_: super_gradients.SgModel
-  experiment_name: ${experiment_name}
-  model_checkpoints_location: ${model_checkpoints_location}
-  ckpt_root_dir: ${ckpt_root_dir}
-  multi_gpu: ${multi_gpu}
-
 architecture: repvgg_a0

@@ -49,14 +49,7 @@ experiment_name: resnet50_imagenet
 ckpt_root_dir:
 
 multi_gpu:
-  _target_: super_gradients.training.sg_model.MultiGPUMode
+  _target_: super_gradients.training.sg_trainer.MultiGPUMode
   value: 'DDP'
 
-sg_model:
-  _target_: super_gradients.SgModel
-  experiment_name: ${experiment_name}
-  model_checkpoints_location: ${model_checkpoints_location}
-  ckpt_root_dir: ${ckpt_root_dir}
-  multi_gpu: ${multi_gpu}
-
 architecture: resnet50

@@ -26,7 +26,7 @@ training_hyperparams:
 
 arch_params:
   teacher_input_adapter:
-    _target_: super_gradients.training.utils.kd_model_utils.NormalizationAdapter
+    _target_: super_gradients.training.utils.kd_trainer_utils.NormalizationAdapter
     mean_original: [0.485, 0.456, 0.406]
     std_original: [0.229, 0.224, 0.225]
     mean_required: [0.5, 0.5, 0.5]
@@ -73,16 +73,9 @@ experiment_name: resnet50_imagenet_KD_Model
 ckpt_root_dir:
 
 multi_gpu:
-  _target_: super_gradients.training.sg_model.MultiGPUMode
+  _target_: super_gradients.training.sg_trainer.MultiGPUMode
   value: DDP
 
-sg_model:
-  _target_: super_gradients.KDModel
-  experiment_name: ${experiment_name}
-  model_checkpoints_location: ${model_checkpoints_location}
-  ckpt_root_dir: ${ckpt_root_dir}
-  multi_gpu: ${multi_gpu}
-
 architecture: kd_module
 student_architecture: resnet50
 teacher_architecture: beit_base_patch16_224

@@ -47,10 +47,4 @@ load_weights_only: True
 
 experiment_name: vit_base_imagenet1k
 
-sg_model:
-  _target_: super_gradients.SgModel
-  experiment_name: ${experiment_name}
-  model_checkpoints_location: ${model_checkpoints_location}
-  multi_gpu: AUTO
-
 architecture: vit_base

@@ -20,11 +20,5 @@ experiment_name: test
 
 model_checkpoints_location: local
 
-sg_model:
-  _target_: super_gradients.SgModel
-  experiment_name: ${experiment_name}
-  model_checkpoints_location: ${model_checkpoints_location}
-  multi_gpu: Off
-
 architecture: resnet18
 

@@ -1,11 +1,14 @@
 # PACKAGE IMPORTS FOR EXTERNAL USAGE
 import super_gradients.training.utils.distributed_training_utils as distributed_training_utils
-from super_gradients.training.datasets import datasets_utils, DataAugmentation, TestDatasetInterface, SegmentationTestDatasetInterface,\
+from super_gradients.training.datasets import datasets_utils, DataAugmentation, TestDatasetInterface, SegmentationTestDatasetInterface, \
     DetectionTestDatasetInterface, ClassificationTestDatasetInterface
 from super_gradients.training.models import ARCHITECTURES
-from super_gradients.training.sg_model import SgModel, MultiGPUMode, StrictLoad
+from super_gradients.training.sg_trainer import Trainer
+from super_gradients.training.kd_trainer import KDTrainer
+from super_gradients.training.sg_model import SgModel
 from super_gradients.training.kd_model import KDModel
+from super_gradients.common import MultiGPUMode, StrictLoad, EvaluationType
 
-__all__ = ['distributed_training_utils', 'datasets_utils', 'DataAugmentation', 'TestDatasetInterface', 'ARCHITECTURES', 'SgModel',
-           'KDModel', 'MultiGPUMode', 'TestDatasetInterface', 'SegmentationTestDatasetInterface', 'DetectionTestDatasetInterface',
-           'ClassificationTestDatasetInterface', 'StrictLoad']
+__all__ = ['distributed_training_utils', 'datasets_utils', 'DataAugmentation', 'TestDatasetInterface',
+           'ARCHITECTURES', 'Trainer', 'KDTrainer', 'MultiGPUMode', 'TestDatasetInterface', 'SegmentationTestDatasetInterface', 'DetectionTestDatasetInterface',
+           'ClassificationTestDatasetInterface', 'StrictLoad', 'SgModel', 'EvaluationType', 'KDModel']

@@ -14,7 +14,7 @@ import torch.distributed as dist
 from super_gradients.common.sg_loggers.abstract_sg_logger import AbstractSGLogger
 
 from super_gradients.common.abstractions.abstract_logger import get_logger
-from deprecated import deprecated
+from deprecate import deprecated
 from matplotlib.patches import Rectangle
 from torchvision.datasets import ImageFolder
 from super_gradients.training.datasets.auto_augment import rand_augment_transform
@@ -75,7 +75,7 @@ def get_mean_and_std_torch(data_dir=None, dataloader=None, num_workers=4, Random
     return mean.view(-1).cpu().numpy().tolist(), std.view(-1).cpu().numpy().tolist()
 
 
-@deprecated(reason='Use get_mean_and_std_torch() instead. It is faster and more accurate')
+@deprecated(target=get_mean_and_std_torch, deprecated_in="2.1.0", remove_in="3.0.0")
 def get_mean_and_std(dataset):
     '''Compute the mean and std value of dataset.'''
     dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=1)

@@ -6,7 +6,7 @@ class KDModelException(Exception):
     """
 
     def __init__(self, desc):
-        self.message = "KDModel: " + desc
+        self.message = "KDTrainer: " + desc
         super().__init__(self.message)
 
 
@@ -48,7 +48,7 @@ class InconsistentParamsException(KDModelException):
 
 
 class UnsupportedKDModelArgException(KDModelException):
-    """Exception raised for unsupported args that might be supported for SgModel but not for KDModel.
+    """Exception raised for unsupported args that might be supported for Trainer but not for KDTrainer.
 
     Attributes:
         message -- explanation of the error
@@ -66,7 +66,7 @@ class TeacherKnowledgeException(KDModelException):
     """
 
     def __init__(self):
-        super().__init__("Expected: at least one of: teacher_pretrained_weights, teacher_checkpoint_path or load_kd_model_checkpoint=True")
+        super().__init__("Expected: at least one of: teacher_pretrained_weights, teacher_checkpoint_path or load_kd_trainer_checkpoint=True")
 
 
 class UndefinedNumClassesException(KDModelException):

@@ -29,4 +29,4 @@ class IllegalDataloaderInitialization(Exception):
 
     def __init__(self):
         super().__init__(
-            "train_loader, valid_loader and class parameters are required when initializing SgModel with data loaders")
+            "train_loader, valid_loader and class parameters are required when initializing Trainer with data loaders")

@@ -1,275 +1,9 @@
-import torch.nn
+from deprecate import deprecated
 
-from super_gradients.training.models.all_architectures import KD_ARCHITECTURES
-from super_gradients.training.models.kd_modules.kd_module import KDModule
-from super_gradients.training.sg_model import SgModel
-from typing import Union
-from super_gradients.common.abstractions.abstract_logger import get_logger
-from super_gradients.training import utils as core_utils
-from super_gradients.training.pretrained_models import PRETRAINED_NUM_CLASSES
-from super_gradients.training.utils import get_param
-from super_gradients.training.utils.checkpoint_utils import read_ckpt_state_dict, \
-    load_checkpoint_to_model
-from super_gradients.training.exceptions.kd_model_exceptions import ArchitectureKwargsException, \
-    UnsupportedKDArchitectureException, InconsistentParamsException, UnsupportedKDModelArgException, \
-    TeacherKnowledgeException, UndefinedNumClassesException
-from super_gradients.training.utils.callbacks import KDModelMetricsUpdateCallback
-from super_gradients.training.utils.ema import KDModelEMA
-logger = get_logger(__name__)
+from super_gradients.training import KDTrainer
 
 
-class KDModel(SgModel):
-    def __init__(self, *args, **kwargs):
-        super(KDModel, self).__init__(*args, **kwargs)
-        self.student_architecture = None
-        self.teacher_architecture = None
-        self.student_arch_params = None
-        self.teacher_arch_params = None
-
-    def build_model(self,
-                    # noqa: C901 - too complex
-                    architecture: Union[str, KDModule] = 'kd_module',
-                    arch_params={}, checkpoint_params={},
-                    *args, **kwargs):
-        """
-        :param architecture: (Union[str, KDModule]) Defines the network's architecture from models/KD_ARCHITECTURES
-         (default='kd_module')
-
-        :param arch_params: (dict) Architecture H.P. e.g.: block, num_blocks, num_classes, etc to be passed to kd
-            architecture class (discarded when architecture is KDModule instance)
-
-        :param checkpoint_params: (dict) A dictionary like object with the following keys/values:
-
-              student_pretrained_weights:   String describing the dataset of the pretrained weights (for example
-              "imagenent") for the student network.
-
-              teacher_pretrained_weights:   String describing the dataset of the pretrained weights (for example
-              "imagenent") for the teacher network.
-
-              teacher_checkpoint_path:    Local path to the teacher's checkpoint. Note that when passing pretrained_weights
-                                   through teacher_arch_params these weights will be overridden by the
-                                   pretrained checkpoint. (default=None)
-
-              load_kd_model_checkpoint:   Whether to load an entire KDModule checkpoint (used to continue KD training)
-               (default=False)
-
-              kd_model_source_ckpt_folder_name: Folder name to load an entire KDModule checkpoint from
-                (self.experiment_name if none is given) to resume KD training (default=None)
-
-              kd_model_external_checkpoint_path: The path to the external checkpoint to be loaded. Can be absolute or relative
-                                               (ie: path/to/checkpoint.pth). If provided, will automatically attempt to
-                                               load the checkpoint even if the load_checkpoint flag is not provided.
-                                               (deafult=None)
-
-        :keyword student_architecture: (Union[str, SgModule]) Defines the student's architecture from
-            models/ALL_ARCHITECTURES (when str), or directly defined the student network (when SgModule).
-
-        :keyword teacher_architecture: (Union[str, SgModule]) Defines the teacher's architecture from
-            models/ALL_ARCHITECTURES (when str), or directly defined the teacher network (when SgModule).
-
-        :keyword student_arch_params: (dict) Architecture H.P. e.g.: block, num_blocks, num_classes, etc for student
-            net. (deafult={})
-
-        :keyword teacher_arch_params: (dict) Architecture H.P. e.g.: block, num_blocks, num_classes, etc for teacher
-            net. (deafult={})
-
-        :keyword run_teacher_on_eval: (bool)- whether to run self.teacher at eval mode regardless of self.train(mode)
-
-
-        """
-        kwargs.setdefault("student_architecture", None)
-        kwargs.setdefault("teacher_architecture", None)
-        kwargs.setdefault("student_arch_params", {})
-        kwargs.setdefault("teacher_arch_params", {})
-        kwargs.setdefault("run_teacher_on_eval", False)
-
-        self._validate_args(arch_params, architecture, checkpoint_params, **kwargs)
-
-        self.student_architecture = kwargs.get("student_architecture")
-        self.teacher_architecture = kwargs.get("teacher_architecture")
-        self.student_arch_params = kwargs.get("student_arch_params")
-        self.teacher_arch_params = kwargs.get("teacher_arch_params")
-
-        super(KDModel, self).build_model(architecture=architecture, arch_params=arch_params,
-                                         checkpoint_params=checkpoint_params, **kwargs)
-
-    def _validate_args(self, arch_params, architecture, checkpoint_params, **kwargs):
-        student_architecture = get_param(kwargs, "student_architecture")
-        teacher_architecture = get_param(kwargs, "teacher_architecture")
-        student_arch_params = get_param(kwargs, "student_arch_params")
-        teacher_arch_params = get_param(kwargs, "teacher_arch_params")
-
-        if get_param(checkpoint_params, 'pretrained_weights') is not None:
-            raise UnsupportedKDModelArgException("pretrained_weights", "checkpoint_params")
-
-        if not isinstance(architecture, KDModule):
-            if student_architecture is None or teacher_architecture is None:
-                raise ArchitectureKwargsException()
-            if architecture not in KD_ARCHITECTURES.keys():
-                raise UnsupportedKDArchitectureException(architecture)
-
-        # DERIVE NUMBER OF CLASSES FROM DATASET INTERFACE IF NOT SPECIFIED OR ARCH PARAMS FOR TEACHER AND STUDENT
-        self._validate_num_classes(student_arch_params, teacher_arch_params)
-
-        arch_params['num_classes'] = student_arch_params['num_classes']
-
-        # MAKE SURE TEACHER'S PRETRAINED NUM CLASSES EQUALS TO THE ONES BELONGING TO STUDENT AS WE CAN'T REPLACE
-        # THE TEACHER'S HEAD
-        teacher_pretrained_weights = core_utils.get_param(checkpoint_params, 'teacher_pretrained_weights',
-                                                          default_val=None)
-        if teacher_pretrained_weights is not None:
-            teacher_pretrained_num_classes = PRETRAINED_NUM_CLASSES[teacher_pretrained_weights]
-            if teacher_pretrained_num_classes != teacher_arch_params['num_classes']:
-                raise InconsistentParamsException("Pretrained dataset number of classes", "teacher's arch params",
-                                                  "number of classes", "student's number of classes")
-
-        teacher_checkpoint_path = get_param(checkpoint_params, "teacher_checkpoint_path")
-        load_kd_model_checkpoint = get_param(checkpoint_params, "load_checkpoint")
-
-        # CHECK THAT TEACHER NETWORK HOLDS KNOWLEDGE FOR THE STUDENT TO LEARN FROM OR THAT WE ARE LOADING AN ENTIRE KD
-        if not (teacher_pretrained_weights or teacher_checkpoint_path or load_kd_model_checkpoint or isinstance(teacher_architecture, torch.nn.Module)):
-            raise TeacherKnowledgeException()
-
-    def _validate_num_classes(self, student_arch_params, teacher_arch_params):
-        """
-        Checks validity of num_classes for num_classes (i.e existence and consistency between subnets)
-
-        :param student_arch_params: (dict) Architecture H.P. e.g.: block, num_blocks, num_classes, etc for student
-        :param teacher_arch_params: (dict) Architecture H.P. e.g.: block, num_blocks, num_classes, etc for teacher
-
-        """
-        self._validate_subnet_num_classes(student_arch_params)
-        self._validate_subnet_num_classes(teacher_arch_params)
-        if teacher_arch_params['num_classes'] != student_arch_params['num_classes']:
-            raise InconsistentParamsException("num_classes", "student_arch_params", "num_classes",
-                                              "teacher_arch_params")
-
-    def _validate_subnet_num_classes(self, subnet_arch_params):
-        """
-        Derives num_classes in student_arch_params/teacher_arch_params from dataset interface or raises an error
-         when none is given
-
-        :param subnet_arch_params: Arch params for student/teacher
-
-        """
-
-        if 'num_classes' not in subnet_arch_params.keys():
-            if self.dataset_interface is None:
-                raise UndefinedNumClassesException()
-            else:
-                subnet_arch_params['num_classes'] = len(self.classes)
-
-    def _instantiate_net(self, architecture: Union[KDModule, KDModule.__class__, str], arch_params: dict,
-                         checkpoint_params: dict, *args, **kwargs) -> tuple:
-        """
-        Instantiates kd_module according to architecture and arch_params, handles pretrained weights for the student
-         and teacher networks, and the required module manipulation (i.e head replacement) for the teacher network.
-
-        :param architecture: String, KDModule or uninstantiated KDModule class describing the netowrks architecture.
-        :param arch_params: Architecture's parameters passed to networks c'tor.
-        :param checkpoint_params: checkpoint loading related parameters dictionary with 'pretrained_weights' key,
-            s.t it's value is a string describing the dataset of the pretrained weights (for example "imagenent").
-
-        :return: instantiated netowrk i.e KDModule, architecture_class (will be none when architecture is not str)
-        """
-
-        student_architecture = get_param(kwargs, "student_architecture")
-        teacher_architecture = get_param(kwargs, "teacher_architecture")
-        student_arch_params = get_param(kwargs, "student_arch_params")
-        teacher_arch_params = get_param(kwargs, "teacher_arch_params")
-        student_arch_params = core_utils.HpmStruct(**student_arch_params)
-        teacher_arch_params = core_utils.HpmStruct(**teacher_arch_params)
-        student_pretrained_weights = get_param(checkpoint_params, 'student_pretrained_weights')
-        teacher_pretrained_weights = get_param(checkpoint_params, 'teacher_pretrained_weights')
-
-        student = super()._instantiate_net(student_architecture, student_arch_params,
-                                           {"pretrained_weights": student_pretrained_weights})
-        teacher = super()._instantiate_net(teacher_architecture, teacher_arch_params,
-                                           {"pretrained_weights": teacher_pretrained_weights})
-
-        run_teacher_on_eval = get_param(kwargs, "run_teacher_on_eval", default_val=False)
-
-        if isinstance(architecture, str):
-            architecture_cls = KD_ARCHITECTURES[architecture]
-            net = architecture_cls(arch_params=arch_params, student=student, teacher=teacher,
-                                   run_teacher_on_eval=run_teacher_on_eval)
-        elif isinstance(architecture, KDModule.__class__):
-            net = architecture(arch_params=arch_params, student=student, teacher=teacher,
-                               run_teacher_on_eval=run_teacher_on_eval)
-        else:
-            net = architecture
-
-        return net
-
-    def _load_checkpoint_to_model(self):
-        """
-        Initializes teacher weights with teacher_checkpoint_path if needed, then handles checkpoint loading for
-         the entire KD network following the same logic as in SgModel.
-        """
-        teacher_checkpoint_path = get_param(self.checkpoint_params, "teacher_checkpoint_path")
-        teacher_net = self.net.module.teacher
-
-        if teacher_checkpoint_path is not None:
-
-            #  WARN THAT TEACHER_CKPT WILL OVERRIDE TEACHER'S PRETRAINED WEIGHTS
-            teacher_pretrained_weights = get_param(self.checkpoint_params, "teacher_pretrained_weights")
-            if teacher_pretrained_weights:
-                logger.warning(
-                    teacher_checkpoint_path + " checkpoint is "
-                                              "overriding " + teacher_pretrained_weights + " for teacher model")
-
-            # ALWAYS LOAD ITS EMA IF IT EXISTS
-            load_teachers_ema = 'ema_net' in read_ckpt_state_dict(teacher_checkpoint_path).keys()
-            load_checkpoint_to_model(ckpt_local_path=teacher_checkpoint_path,
-                                     load_backbone=False,
-                                     net=teacher_net,
-                                     strict='no_key_matching',
-                                     load_weights_only=True,
-                                     load_ema_as_net=load_teachers_ema)
-
-        super(KDModel, self)._load_checkpoint_to_model()
-
-    def _add_metrics_update_callback(self, phase):
-        """
-        Adds KDModelMetricsUpdateCallback to be fired at phase
-
-        :param phase: Phase for the metrics callback to be fired at
-        """
-        self.phase_callbacks.append(KDModelMetricsUpdateCallback(phase))
-
-    def _get_hyper_param_config(self):
-        """
-        Creates a training hyper param config for logging with additional KD related hyper params.
-        """
-        hyper_param_config = super()._get_hyper_param_config()
-        hyper_param_config.update({"student_architecture": self.student_architecture,
-                                   "teacher_architecture": self.teacher_architecture,
-                                   "student_arch_params": self.student_arch_params,
-                                   "teacher_arch_params": self.teacher_arch_params
-                                   })
-        return hyper_param_config
-
-    def _instantiate_ema_model(self, decay: float = 0.9999, beta: float = 15, exp_activation: bool = True) -> KDModelEMA:
-        """Instantiate KD ema model for KDModule.
-
-        If the model is of class KDModule, the instance will be adapted to work on knowledge distillation.
-        :param decay:           the maximum decay value. as the training process advances, the decay will climb towards
-                                this value until the EMA_t+1 = EMA_t * decay + TRAINING_MODEL * (1- decay)
-        :param beta:            the exponent coefficient. The higher the beta, the sooner in the training the decay will
-                                saturate to its final value. beta=15 is ~40% of the training process.
-        :param exp_activation:
-        """
-        return KDModelEMA(self.net, decay, beta, exp_activation)
-
-    def _save_best_checkpoint(self, epoch, state):
-        """
-        Overrides parent best_ckpt saving to modify the state dict so that we only save the student.
-        """
-        if self.ema:
-            best_net = core_utils.WrappedModel(self.ema_model.ema.module.student)
-            state.pop("ema_net")
-        else:
-            best_net = core_utils.WrappedModel(self.net.module.student)
-
-        state["net"] = best_net.state_dict()
-        self.sg_logger.add_checkpoint(tag=self.ckpt_best_name, state_dict=state, global_step=epoch)
+@deprecated(target=KDTrainer, deprecated_in='2.3.0', remove_in='3.0.0')
+class KDModel(KDTrainer):
+    def __init__(self, experiment_name: str, *args, **kwargs):
+        super().__init__(experiment_name, *args, **kwargs)

from super_gradients.training.trainer import Trainer


class KDTrainer(Trainer):
    """
    Class for running SuperGradient's recipes for KD Models.
    See train_from_kd_recipe example in the examples directory to demonstrate it's usage.
    """

    @classmethod
    def build_model(cls, cfg):
        cfg.sg_model.build_model(student_architecture=cfg.student_architecture,
                                 teacher_architecture=cfg.teacher_architecture,
                                 arch_params=cfg.arch_params, student_arch_params=cfg.student_arch_params,
                                 teacher_arch_params=cfg.teacher_arch_params,
                                 checkpoint_params=cfg.checkpoint_params, run_teacher_on_eval=cfg.run_teacher_on_eval)

# PACKAGE IMPORTS FOR EXTERNAL USAGE

from super_gradients.training.kd_trainer.kd_trainer import KDTrainer

__all__ = ['KDTrainer']

import hydra
import torch.nn
from omegaconf import DictConfig
from torch.utils.data import DataLoader

from super_gradients.common import MultiGPUMode
from super_gradients.training.models.all_architectures import KD_ARCHITECTURES
from super_gradients.training.models.kd_modules.kd_module import KDModule
from super_gradients.training.sg_trainer import Trainer
from typing import Union, List, Any
from super_gradients.common.abstractions.abstract_logger import get_logger
from super_gradients.training import utils as core_utils
from super_gradients.training.pretrained_models import PRETRAINED_NUM_CLASSES
from super_gradients.training.utils import get_param
from super_gradients.training.utils.checkpoint_utils import read_ckpt_state_dict, \
    load_checkpoint_to_model
from super_gradients.training.exceptions.kd_trainer_exceptions import ArchitectureKwargsException, \
    UnsupportedKDArchitectureException, InconsistentParamsException, UnsupportedKDModelArgException, \
    TeacherKnowledgeException, UndefinedNumClassesException
from super_gradients.training.utils.callbacks import KDModelMetricsUpdateCallback
from super_gradients.training.utils.detection_utils import DetectionPostPredictionCallback
from super_gradients.training.utils.ema import KDModelEMA
from super_gradients.training.utils.sg_trainer_utils import parse_args

logger = get_logger(__name__)


class KDTrainer(Trainer):
    def __init__(self, experiment_name: str, device: str = None, multi_gpu: Union[MultiGPUMode, str] = MultiGPUMode.OFF,
                 model_checkpoints_location: str = 'local', overwrite_local_checkpoint: bool = True, ckpt_name: str = 'ckpt_latest.pth',
                 post_prediction_callback: DetectionPostPredictionCallback = None, ckpt_root_dir: str = None, train_loader: DataLoader = None,
                 valid_loader: DataLoader = None, test_loader: DataLoader = None, classes: List[Any] = None):

        super().__init__(experiment_name, device, multi_gpu, model_checkpoints_location, overwrite_local_checkpoint, ckpt_name, post_prediction_callback,
                         ckpt_root_dir, train_loader, valid_loader, test_loader, classes)
        self.student_architecture = None
        self.teacher_architecture = None
        self.student_arch_params = None
        self.teacher_arch_params = None

    @classmethod
    def train_from_config(cls, cfg: Union[DictConfig, dict]) -> None:
        """
        Trains according to cfg recipe configuration.

        @param cfg: The parsed DictConfig from yaml recipe files
        @return: output of kd_trainer.train(...) (i.e results tuple)
        """
        # INSTANTIATE ALL OBJECTS IN CFG
        cfg = hydra.utils.instantiate(cfg)

        kwargs = parse_args(cfg, cls.__init__)

        trainer = KDTrainer(**kwargs)

        # CONNECT THE DATASET INTERFACE WITH DECI MODEL
        trainer.connect_dataset_interface(cfg.dataset_interface, data_loader_num_workers=cfg.data_loader_num_workers)

        # BUILD NETWORK
        trainer.build_model(student_architecture=cfg.student_architecture,
                            teacher_architecture=cfg.teacher_architecture,
                            arch_params=cfg.arch_params, student_arch_params=cfg.student_arch_params,
                            teacher_arch_params=cfg.teacher_arch_params,
                            checkpoint_params=cfg.checkpoint_params, run_teacher_on_eval=cfg.run_teacher_on_eval)

        # TRAIN
        trainer.train(training_params=cfg.training_hyperparams)

    def build_model(self,
                    # noqa: C901 - too complex
                    architecture: Union[str, KDModule] = 'kd_module',
                    arch_params={}, checkpoint_params={},
                    *args, **kwargs):
        """
        :param architecture: (Union[str, KDModule]) Defines the network's architecture from models/KD_ARCHITECTURES
         (default='kd_module')

        :param arch_params: (dict) Architecture H.P. e.g.: block, num_blocks, num_classes, etc to be passed to kd
            architecture class (discarded when architecture is KDModule instance)

        :param checkpoint_params: (dict) A dictionary like object with the following keys/values:

              student_pretrained_weights:   String describing the dataset of the pretrained weights (for example
              "imagenent") for the student network.

              teacher_pretrained_weights:   String describing the dataset of the pretrained weights (for example
              "imagenent") for the teacher network.

              teacher_checkpoint_path:    Local path to the teacher's checkpoint. Note that when passing pretrained_weights
                                   through teacher_arch_params these weights will be overridden by the
                                   pretrained checkpoint. (default=None)

              load_kd_model_checkpoint:   Whether to load an entire KDModule checkpoint (used to continue KD training)
               (default=False)

              kd_model_source_ckpt_folder_name: Folder name to load an entire KDModule checkpoint from
                (self.experiment_name if none is given) to resume KD training (default=None)

              kd_model_external_checkpoint_path: The path to the external checkpoint to be loaded. Can be absolute or relative
                                               (ie: path/to/checkpoint.pth). If provided, will automatically attempt to
                                               load the checkpoint even if the load_checkpoint flag is not provided.
                                               (deafult=None)

        :keyword student_architecture: (Union[str, SgModule]) Defines the student's architecture from
            models/ALL_ARCHITECTURES (when str), or directly defined the student network (when SgModule).

        :keyword teacher_architecture: (Union[str, SgModule]) Defines the teacher's architecture from
            models/ALL_ARCHITECTURES (when str), or directly defined the teacher network (when SgModule).

        :keyword student_arch_params: (dict) Architecture H.P. e.g.: block, num_blocks, num_classes, etc for student
            net. (deafult={})

        :keyword teacher_arch_params: (dict) Architecture H.P. e.g.: block, num_blocks, num_classes, etc for teacher
            net. (deafult={})

        :keyword run_teacher_on_eval: (bool)- whether to run self.teacher at eval mode regardless of self.train(mode)


        """
        kwargs.setdefault("student_architecture", None)
        kwargs.setdefault("teacher_architecture", None)
        kwargs.setdefault("student_arch_params", {})
        kwargs.setdefault("teacher_arch_params", {})
        kwargs.setdefault("run_teacher_on_eval", False)

        self._validate_args(arch_params, architecture, checkpoint_params, **kwargs)

        self.student_architecture = kwargs.get("student_architecture")
        self.teacher_architecture = kwargs.get("teacher_architecture")
        self.student_arch_params = kwargs.get("student_arch_params")
        self.teacher_arch_params = kwargs.get("teacher_arch_params")

        super(KDTrainer, self).build_model(architecture=architecture, arch_params=arch_params,
                                           checkpoint_params=checkpoint_params, **kwargs)

    def _validate_args(self, arch_params, architecture, checkpoint_params, **kwargs):
        student_architecture = get_param(kwargs, "student_architecture")
        teacher_architecture = get_param(kwargs, "teacher_architecture")
        student_arch_params = get_param(kwargs, "student_arch_params")
        teacher_arch_params = get_param(kwargs, "teacher_arch_params")

        if get_param(checkpoint_params, 'pretrained_weights') is not None:
            raise UnsupportedKDModelArgException("pretrained_weights", "checkpoint_params")

        if not isinstance(architecture, KDModule):
            if student_architecture is None or teacher_architecture is None:
                raise ArchitectureKwargsException()
            if architecture not in KD_ARCHITECTURES.keys():
                raise UnsupportedKDArchitectureException(architecture)

        # DERIVE NUMBER OF CLASSES FROM DATASET INTERFACE IF NOT SPECIFIED OR ARCH PARAMS FOR TEACHER AND STUDENT
        self._validate_num_classes(student_arch_params, teacher_arch_params)

        arch_params['num_classes'] = student_arch_params['num_classes']

        # MAKE SURE TEACHER'S PRETRAINED NUM CLASSES EQUALS TO THE ONES BELONGING TO STUDENT AS WE CAN'T REPLACE
        # THE TEACHER'S HEAD
        teacher_pretrained_weights = core_utils.get_param(checkpoint_params, 'teacher_pretrained_weights',
                                                          default_val=None)
        if teacher_pretrained_weights is not None:
            teacher_pretrained_num_classes = PRETRAINED_NUM_CLASSES[teacher_pretrained_weights]
            if teacher_pretrained_num_classes != teacher_arch_params['num_classes']:
                raise InconsistentParamsException("Pretrained dataset number of classes", "teacher's arch params",
                                                  "number of classes", "student's number of classes")

        teacher_checkpoint_path = get_param(checkpoint_params, "teacher_checkpoint_path")
        load_kd_model_checkpoint = get_param(checkpoint_params, "load_checkpoint")

        # CHECK THAT TEACHER NETWORK HOLDS KNOWLEDGE FOR THE STUDENT TO LEARN FROM OR THAT WE ARE LOADING AN ENTIRE KD
        if not (teacher_pretrained_weights or teacher_checkpoint_path or load_kd_model_checkpoint or isinstance(teacher_architecture, torch.nn.Module)):
            raise TeacherKnowledgeException()

    def _validate_num_classes(self, student_arch_params, teacher_arch_params):
        """
        Checks validity of num_classes for num_classes (i.e existence and consistency between subnets)

        :param student_arch_params: (dict) Architecture H.P. e.g.: block, num_blocks, num_classes, etc for student
        :param teacher_arch_params: (dict) Architecture H.P. e.g.: block, num_blocks, num_classes, etc for teacher

        """
        self._validate_subnet_num_classes(student_arch_params)
        self._validate_subnet_num_classes(teacher_arch_params)
        if teacher_arch_params['num_classes'] != student_arch_params['num_classes']:
            raise InconsistentParamsException("num_classes", "student_arch_params", "num_classes",
                                              "teacher_arch_params")

    def _validate_subnet_num_classes(self, subnet_arch_params):
        """
        Derives num_classes in student_arch_params/teacher_arch_params from dataset interface or raises an error
         when none is given

        :param subnet_arch_params: Arch params for student/teacher

        """

        if 'num_classes' not in subnet_arch_params.keys():
            if self.dataset_interface is None:
                raise UndefinedNumClassesException()
            else:
                subnet_arch_params['num_classes'] = len(self.classes)

    def _instantiate_net(self, architecture: Union[KDModule, KDModule.__class__, str], arch_params: dict,
                         checkpoint_params: dict, *args, **kwargs) -> tuple:
        """
        Instantiates kd_module according to architecture and arch_params, handles pretrained weights for the student
         and teacher networks, and the required module manipulation (i.e head replacement) for the teacher network.

        :param architecture: String, KDModule or uninstantiated KDModule class describing the netowrks architecture.
        :param arch_params: Architecture's parameters passed to networks c'tor.
        :param checkpoint_params: checkpoint loading related parameters dictionary with 'pretrained_weights' key,
            s.t it's value is a string describing the dataset of the pretrained weights (for example "imagenent").

        :return: instantiated netowrk i.e KDModule, architecture_class (will be none when architecture is not str)
        """

        student_architecture = get_param(kwargs, "student_architecture")
        teacher_architecture = get_param(kwargs, "teacher_architecture")
        student_arch_params = get_param(kwargs, "student_arch_params")
        teacher_arch_params = get_param(kwargs, "teacher_arch_params")
        student_arch_params = core_utils.HpmStruct(**student_arch_params)
        teacher_arch_params = core_utils.HpmStruct(**teacher_arch_params)
        student_pretrained_weights = get_param(checkpoint_params, 'student_pretrained_weights')
        teacher_pretrained_weights = get_param(checkpoint_params, 'teacher_pretrained_weights')

        student = super()._instantiate_net(student_architecture, student_arch_params,
                                           {"pretrained_weights": student_pretrained_weights})
        teacher = super()._instantiate_net(teacher_architecture, teacher_arch_params,
                                           {"pretrained_weights": teacher_pretrained_weights})

        run_teacher_on_eval = get_param(kwargs, "run_teacher_on_eval", default_val=False)

        if isinstance(architecture, str):
            architecture_cls = KD_ARCHITECTURES[architecture]
            net = architecture_cls(arch_params=arch_params, student=student, teacher=teacher,
                                   run_teacher_on_eval=run_teacher_on_eval)
        elif isinstance(architecture, KDModule.__class__):
            net = architecture(arch_params=arch_params, student=student, teacher=teacher,
                               run_teacher_on_eval=run_teacher_on_eval)
        else:
            net = architecture

        return net

    def _load_checkpoint_to_model(self):
        """
        Initializes teacher weights with teacher_checkpoint_path if needed, then handles checkpoint loading for
         the entire KD network following the same logic as in Trainer.
        """
        teacher_checkpoint_path = get_param(self.checkpoint_params, "teacher_checkpoint_path")
        teacher_net = self.net.module.teacher

        if teacher_checkpoint_path is not None:

            #  WARN THAT TEACHER_CKPT WILL OVERRIDE TEACHER'S PRETRAINED WEIGHTS
            teacher_pretrained_weights = get_param(self.checkpoint_params, "teacher_pretrained_weights")
            if teacher_pretrained_weights:
                logger.warning(
                    teacher_checkpoint_path + " checkpoint is "
                                              "overriding " + teacher_pretrained_weights + " for teacher model")

            # ALWAYS LOAD ITS EMA IF IT EXISTS
            load_teachers_ema = 'ema_net' in read_ckpt_state_dict(teacher_checkpoint_path).keys()
            load_checkpoint_to_model(ckpt_local_path=teacher_checkpoint_path,
                                     load_backbone=False,
                                     net=teacher_net,
                                     strict='no_key_matching',
                                     load_weights_only=True,
                                     load_ema_as_net=load_teachers_ema)

        super(KDTrainer, self)._load_checkpoint_to_model()

    def _add_metrics_update_callback(self, phase):
        """
        Adds KDModelMetricsUpdateCallback to be fired at phase

        :param phase: Phase for the metrics callback to be fired at
        """
        self.phase_callbacks.append(KDModelMetricsUpdateCallback(phase))

    def _get_hyper_param_config(self):
        """
        Creates a training hyper param config for logging with additional KD related hyper params.
        """
        hyper_param_config = super()._get_hyper_param_config()
        hyper_param_config.update({"student_architecture": self.student_architecture,
                                   "teacher_architecture": self.teacher_architecture,
                                   "student_arch_params": self.student_arch_params,
                                   "teacher_arch_params": self.teacher_arch_params
                                   })
        return hyper_param_config

    def _instantiate_ema_model(self, decay: float = 0.9999, beta: float = 15, exp_activation: bool = True) -> KDModelEMA:
        """Instantiate KD ema model for KDModule.

        If the model is of class KDModule, the instance will be adapted to work on knowledge distillation.
        :param decay:           the maximum decay value. as the training process advances, the decay will climb towards
                                this value until the EMA_t+1 = EMA_t * decay + TRAINING_MODEL * (1- decay)
        :param beta:            the exponent coefficient. The higher the beta, the sooner in the training the decay will
                                saturate to its final value. beta=15 is ~40% of the training process.
        :param exp_activation:
        """
        return KDModelEMA(self.net, decay, beta, exp_activation)

    def _save_best_checkpoint(self, epoch, state):
        """
        Overrides parent best_ckpt saving to modify the state dict so that we only save the student.
        """
        if self.ema:
            best_net = core_utils.WrappedModel(self.ema_model.ema.module.student)
            state.pop("ema_net")
        else:
            best_net = core_utils.WrappedModel(self.net.module.student)

        state["net"] = best_net.state_dict()
        self.sg_logger.add_checkpoint(tag=self.ckpt_best_name, state_dict=state, global_step=epoch)

@@ -480,7 +480,7 @@ class STDCSegmentationBase(SgModule):
     @property
     def backbone(self):
         """
-        For SgModel load_backbone compatibility.
+        For Trainer load_backbone compatibility.
         """
         return self.cp.backbone
 

@@ -1,5 +1,5 @@
 # PACKAGE IMPORTS FOR EXTERNAL USAGE
+from super_gradients.common import MultiGPUMode, StrictLoad, EvaluationType
+from super_gradients.training.sg_model.sg_model import SgModel
 
-from super_gradients.training.sg_model.sg_model import SgModel, MultiGPUMode, StrictLoad
-
-__all__ = ['SgModel', 'MultiGPUMode', 'StrictLoad']
+__all__ = ['SgModel', 'MultiGPUMode', 'StrictLoad', 'EvaluationType']

@@ -1,1826 +1,9 @@
-import inspect
-import os
-import sys
-from copy import deepcopy
-from typing import Union, Tuple, Mapping, List, Any
+from deprecate import deprecated
 
-import numpy as np
-import pkg_resources
-import torch
-import torchvision.transforms as transforms
-from deprecated import deprecated
-from torch import nn
-from torch.utils.data import DataLoader, DistributedSampler
-from torch.cuda.amp import GradScaler, autocast
-from torchmetrics import MetricCollection
-from tqdm import tqdm
-from piptools.scripts.sync import _get_installed_distributions
+from super_gradients.training import Trainer
 
-from super_gradients.common.factories.callbacks_factory import CallbacksFactory
-from super_gradients.training.models.all_architectures import ARCHITECTURES
-from super_gradients.common.decorators.factory_decorator import resolve_param
-from super_gradients.common.environment import env_helpers
-from super_gradients.common.abstractions.abstract_logger import get_logger
-from super_gradients.common.factories.datasets_factory import DatasetsFactory
-from super_gradients.common.factories.list_factory import ListFactory
-from super_gradients.common.factories.losses_factory import LossesFactory
-from super_gradients.common.factories.metrics_factory import MetricsFactory
-from super_gradients.common.sg_loggers import SG_LOGGERS
-from super_gradients.common.sg_loggers.abstract_sg_logger import AbstractSGLogger
-from super_gradients.common.sg_loggers.base_sg_logger import BaseSGLogger
-from super_gradients.training import utils as core_utils
-from super_gradients.training.models import SgModule
-from super_gradients.training.pretrained_models import PRETRAINED_NUM_CLASSES
-from super_gradients.training.utils import sg_model_utils
-from super_gradients.training.utils.quantization_utils import QATCallback
-from super_gradients.training.utils.sg_model_utils import MonitoredValue
-from super_gradients.training import metrics
-from super_gradients.training.exceptions.sg_model_exceptions import UnsupportedOptimizerFormat, \
-    IllegalDataloaderInitialization
-from super_gradients.training.datasets import DatasetInterface
-from super_gradients.training.losses import LOSSES
-from super_gradients.training.metrics.metric_utils import get_metrics_titles, get_metrics_results_tuple, \
-    get_logging_values, \
-    get_metrics_dict, get_train_loop_description_dict
-from super_gradients.training.params import TrainingParams
-from super_gradients.training.utils.detection_utils import DetectionPostPredictionCallback
-from super_gradients.training.utils.distributed_training_utils import MultiGPUModeAutocastWrapper, \
-    reduce_results_tuple_for_ddp, compute_precise_bn_stats
-from super_gradients.training.utils.ema import ModelEMA
-from super_gradients.training.utils.optimizer_utils import build_optimizer
-from super_gradients.training.utils.weight_averaging_utils import ModelWeightAveraging
-from super_gradients.training.metrics import Accuracy, Top5
-from super_gradients.training.utils import random_seed
-from super_gradients.training.utils.checkpoint_utils import get_ckpt_local_path, read_ckpt_state_dict, \
-    load_checkpoint_to_model, load_pretrained_weights
-from super_gradients.training.datasets.datasets_utils import DatasetStatisticsTensorboardLogger
-from super_gradients.training.utils.callbacks import CallbackHandler, Phase, LR_SCHEDULERS_CLS_DICT, PhaseContext, \
-    MetricsUpdateCallback, LR_WARMUP_CLS_DICT, ContextSgMethods, LRCallbackBase
-from super_gradients.common.environment import environment_config
-from super_gradients.training.utils import HpmStruct
-from super_gradients.training.datasets.samplers.infinite_sampler import InfiniteSampler
 
-from super_gradients.common import StrictLoad, MultiGPUMode, EvaluationType
-logger = get_logger(__name__)
-
-
-class SgModel:
-    """
-    SuperGradient Model - Base Class for Sg Models
-
-    Methods
-    -------
-    train(max_epochs : int, initial_epoch : int, save_model : bool)
-        the main function used for the training, h.p. updating, logging etc.
-
-    predict(idx : int)
-        returns the predictions and label of the current inputs
-
-    test(epoch : int, idx : int, save : bool):
-        returns the test loss, accuracy and runtime
-    """
-
-    def __init__(self, experiment_name: str, device: str = None, multi_gpu: Union[MultiGPUMode, str] = MultiGPUMode.OFF,
-                 model_checkpoints_location: str = 'local',
-                 overwrite_local_checkpoint: bool = True, ckpt_name: str = 'ckpt_latest.pth',
-                 post_prediction_callback: DetectionPostPredictionCallback = None, ckpt_root_dir: str = None,
-                 train_loader: DataLoader = None, valid_loader: DataLoader = None, test_loader: DataLoader = None,
-                 classes: List[Any] = None):
-        """
-
-        :param experiment_name:                      Used for logging and loading purposes
-        :param device:                          If equal to 'cpu' runs on the CPU otherwise on GPU
-        :param multi_gpu:                       If True, runs on all available devices
-        :param model_checkpoints_location:      If set to 's3' saves the Checkpoints in AWS S3
-                                                otherwise saves the Checkpoints Locally
-        :param overwrite_local_checkpoint:      If set to False keeps the current local checkpoint when importing
-                                                checkpoint from cloud service, otherwise overwrites the local checkpoints file
-        :param ckpt_name:                       The Checkpoint to Load
-        :param ckpt_root_dir:                   Local root directory path where all experiment logging directories will
-                                                reside. When none is give, it is assumed that
-                                                pkg_resources.resource_filename('checkpoints', "") exists and will be used.
-        :param train_loader:                    Training set Dataloader instead of using DatasetInterface, must pass "valid_loader"
-                                                and "classes" along with it
-        :param valid_loader:                    Validation set Dataloader
-        :param test_loader:                     Test set Dataloader
-        :param classes:                         List of class labels
-
-        """
-        # SET THE EMPTY PROPERTIES
-        self.net, self.architecture, self.arch_params, self.dataset_interface = None, None, None, None
-        self.device, self.multi_gpu = None, None
-        self.ema = None
-        self.ema_model = None
-        self.sg_logger = None
-        self.update_param_groups = None
-        self.post_prediction_callback = None
-        self.criterion = None
-        self.training_params = None
-        self.scaler = None
-        self.phase_callbacks = None
-        self.checkpoint_params = None
-        self.pre_prediction_callback = None
-
-        # SET THE DEFAULT PROPERTIES
-        self.half_precision = False
-        self.load_checkpoint = False
-        self.load_backbone = False
-        self.load_weights_only = False
-        self.ddp_silent_mode = False
-        self.source_ckpt_folder_name = None
-        self.model_weight_averaging = None
-        self.average_model_checkpoint_filename = 'average_model.pth'
-        self.start_epoch = 0
-        self.best_metric = np.inf
-        self.external_checkpoint_path = None
-        self.strict_load = StrictLoad.ON
-        self.load_ema_as_net = False
-        self.ckpt_best_name = 'ckpt_best.pth'
-        self.enable_qat = False
-        self.qat_params = {}
-        self._infinite_train_loader = False
-
-        # DETERMINE THE LOCATION OF THE LOSS AND ACCURACY IN THE RESULTS TUPLE OUTPUTED BY THE TEST
-        self.loss_idx_in_results_tuple, self.acc_idx_in_results_tuple = None, None
-
-        # METRICS
-        self.loss_logging_items_names = None
-        self.train_metrics = None
-        self.valid_metrics = None
-        self.greater_metric_to_watch_is_better = None
-
-        # SETTING THE PROPERTIES FROM THE CONSTRUCTOR
-        self.experiment_name = experiment_name
-        self.ckpt_name = ckpt_name
-        self.overwrite_local_checkpoint = overwrite_local_checkpoint
-        self.model_checkpoints_location = model_checkpoints_location
-        self._set_dataset_properties(classes, test_loader, train_loader, valid_loader)
-
-        # CREATING THE LOGGING DIR BASED ON THE INPUT PARAMS TO PREVENT OVERWRITE OF LOCAL VERSION
-        if ckpt_root_dir:
-            self.checkpoints_dir_path = os.path.join(ckpt_root_dir, self.experiment_name)
-        elif pkg_resources.resource_exists("checkpoints", ""):
-            self.checkpoints_dir_path = pkg_resources.resource_filename('checkpoints', self.experiment_name)
-        else:
-            raise ValueError("Illegal checkpoints directory: pass ckpt_root_dir that exists, or add 'checkpoints' to"
-                             "resources.")
-
-        # INITIALIZE THE DEVICE FOR THE MODEL
-        self._initialize_device(requested_device=device, requested_multi_gpu=multi_gpu)
-
-        self.post_prediction_callback = post_prediction_callback
-        # SET THE DEFAULTS
-        # TODO: SET DEFAULT TRAINING PARAMS FOR EACH TASK
-
-        default_results_titles = ['Train Loss', 'Train Acc', 'Train Top5', 'Valid Loss', 'Valid Acc', 'Valid Top5']
-
-        self.results_titles = default_results_titles
-
-        self.loss_idx_in_results_tuple, self.acc_idx_in_results_tuple = 0, 1
-        default_train_metrics, default_valid_metrics = MetricCollection([Accuracy(), Top5()]), MetricCollection(
-            [Accuracy(), Top5()])
-
-        default_loss_logging_items_names = ["Loss"]
-
-        self.train_metrics, self.valid_metrics = default_train_metrics, default_valid_metrics
-        self.loss_logging_items_names = default_loss_logging_items_names
-
-        self.train_monitored_values = {}
-        self.valid_monitored_values = {}
-
-    def _set_dataset_properties(self, classes, test_loader, train_loader, valid_loader):
-        if any([train_loader, valid_loader, classes]) and not all([train_loader, valid_loader, classes]):
-            raise IllegalDataloaderInitialization()
-
-        dataset_params = {"batch_size": train_loader.batch_size if train_loader else None,
-                          "val_batch_size": valid_loader.batch_size if valid_loader else None,
-                          "test_batch_size": test_loader.batch_size if test_loader else None,
-                          "dataset_dir": None,
-                          "s3_link": None}
-
-        if train_loader and self.multi_gpu == MultiGPUMode.DISTRIBUTED_DATA_PARALLEL:
-            if not all([isinstance(train_loader.sampler, DistributedSampler),
-                        isinstance(valid_loader.sampler, DistributedSampler),
-                        test_loader is None or isinstance(test_loader.sampler, DistributedSampler)]):
-                logger.warning("DDP training was selected but the dataloader samplers are not of type DistributedSamplers")
-
-        self.dataset_params, self.train_loader, self.valid_loader, self.test_loader, self.classes = \
-            HpmStruct(**dataset_params), train_loader, valid_loader, test_loader, classes
-
-    @resolve_param('dataset_interface', DatasetsFactory())
-    def connect_dataset_interface(self, dataset_interface: DatasetInterface, data_loader_num_workers: int = 8):
-        """
-        :param dataset_interface: DatasetInterface object
-        :param data_loader_num_workers: The number of threads to initialize the Data Loaders with
-            The dataset to be connected
-        """
-        if self.train_loader:
-            logger.warning("Overriding the dataloaders that SgModel was initialized with")
-        self.dataset_interface = dataset_interface
-        self.train_loader, self.valid_loader, self.test_loader, self.classes = \
-            self.dataset_interface.get_data_loaders(batch_size_factor=self.num_devices,
-                                                    num_workers=data_loader_num_workers,
-                                                    distributed_sampler=self.multi_gpu == MultiGPUMode.DISTRIBUTED_DATA_PARALLEL)
-
-        self.dataset_params = self.dataset_interface.get_dataset_params()
-
-    # FIXME - we need to resolve flake8's 'function is too complex' for this function
-    def build_model(self,  # noqa: C901 - too complex
-                    architecture: Union[str, nn.Module],
-                    arch_params={}, checkpoint_params={}, *args, **kwargs):
-        """
-        :param architecture:               Defines the network's architecture from models/ALL_ARCHITECTURES
-        :param arch_params:                Architecture H.P. e.g.: block, num_blocks, num_classes, etc.
-        :param checkpoint_params:          Dictionary like object with the following key:values:
-
-            load_checkpoint:            Load a pre-trained checkpoint
-            strict_load:                See StrictLoad class documentation for details.
-            source_ckpt_folder_name:    folder name to load the checkpoint from (self.experiment_name if none is given)
-            load_weights_only:          loads only the weight from the checkpoint and zeroize the training params
-            load_backbone:              loads the provided checkpoint to self.net.backbone instead of self.net
-            external_checkpoint_path:   The path to the external checkpoint to be loaded. Can be absolute or relative
-                                               (ie: path/to/checkpoint.pth). If provided, will automatically attempt to
-                                               load the checkpoint even if the load_checkpoint flag is not provided.
-
-        """
-        if 'num_classes' not in arch_params.keys():
-            if self.classes is None and self.dataset_interface is None:
-                raise Exception('Error', 'Number of classes not defined in arch params and dataset is not defined')
-            else:
-                arch_params['num_classes'] = len(self.classes)
-
-        self.arch_params = core_utils.HpmStruct(**arch_params)
-        self.checkpoint_params = core_utils.HpmStruct(**checkpoint_params)
-
-        self.net = self._instantiate_net(architecture, self.arch_params, checkpoint_params, *args, **kwargs)
-
-        # SAVE THE ARCHITECTURE FOR NEURAL ARCHITECTURE SEARCH
-
-        self.architecture = architecture
-
-        self._net_to_device()
-
-        # SET THE FLAG FOR DIFFERENT PARAMETER GROUP OPTIMIZER UPDATE
-        self.update_param_groups = hasattr(self.net.module, 'update_param_groups')
-
-        self._load_checkpoint_to_model()
-
-    def _set_ckpt_loading_attributes(self):
-        """
-        Sets checkpoint loading related attributes according to self.checkpoint_params
-        """
-        self.checkpoint = {}
-        self.strict_load = core_utils.get_param(self.checkpoint_params, 'strict_load', default_val=StrictLoad.ON)
-        self.load_ema_as_net = core_utils.get_param(self.checkpoint_params, 'load_ema_as_net', default_val=False)
-        self.source_ckpt_folder_name = core_utils.get_param(self.checkpoint_params, 'source_ckpt_folder_name')
-        self.load_checkpoint = core_utils.get_param(self.checkpoint_params, 'load_checkpoint', default_val=False)
-        self.load_backbone = core_utils.get_param(self.checkpoint_params, 'load_backbone', default_val=False)
-        self.external_checkpoint_path = core_utils.get_param(self.checkpoint_params, 'external_checkpoint_path')
-        if self.load_checkpoint or self.external_checkpoint_path:
-            self.load_weights_only = core_utils.get_param(self.checkpoint_params, 'load_weights_only',
-                                                          default_val=False)
-        self.ckpt_name = core_utils.get_param(self.checkpoint_params, 'ckpt_name', default_val=self.ckpt_name)
-
-    def _net_to_device(self):
-        """
-        Manipulates self.net according to self.multi_gpu
-        """
-        self.net.to(self.device)
-
-        # FOR MULTI-GPU TRAINING (not distributed)
-        self.arch_params.sync_bn = core_utils.get_param(self.arch_params, 'sync_bn', default_val=False)
-        if self.multi_gpu == MultiGPUMode.DATA_PARALLEL:
-            self.net = torch.nn.DataParallel(self.net, device_ids=self.device_ids)
-        elif self.multi_gpu == MultiGPUMode.DISTRIBUTED_DATA_PARALLEL:
-            if self.arch_params.sync_bn:
-                if not self.ddp_silent_mode:
-                    logger.info('DDP - Using Sync Batch Norm... Training time will be affected accordingly')
-                self.net = torch.nn.SyncBatchNorm.convert_sync_batchnorm(self.net).to(self.device)
-
-            local_rank = int(self.device.split(':')[1])
-            self.net = torch.nn.parallel.DistributedDataParallel(self.net,
-                                                                 device_ids=[local_rank],
-                                                                 output_device=local_rank,
-                                                                 find_unused_parameters=True)
-
-        else:
-            self.net = core_utils.WrappedModel(self.net)
-
-    def _train_epoch(self, epoch: int, silent_mode: bool = False) -> tuple:
-        """
-        train_epoch - A single epoch training procedure
-            :param optimizer:   The optimizer for the network
-            :param epoch:       The current epoch
-            :param silent_mode: No verbosity
-        """
-        # SET THE MODEL IN training STATE
-        self.net.train()
-        # THE DISABLE FLAG CONTROLS WHETHER THE PROGRESS BAR IS SILENT OR PRINTS THE LOGS
-        progress_bar_train_loader = tqdm(self.train_loader, bar_format="{l_bar}{bar:10}{r_bar}", dynamic_ncols=True,
-                                         disable=silent_mode)
-        progress_bar_train_loader.set_description(f"Train epoch {epoch}")
-
-        # RESET/INIT THE METRIC LOGGERS
-        self._reset_metrics()
-
-        self.train_metrics.to(self.device)
-        loss_avg_meter = core_utils.utils.AverageMeter()
-
-        context = PhaseContext(epoch=epoch,
-                               optimizer=self.optimizer,
-                               metrics_compute_fn=self.train_metrics,
-                               loss_avg_meter=loss_avg_meter,
-                               criterion=self.criterion,
-                               device=self.device,
-                               lr_warmup_epochs=self.training_params.lr_warmup_epochs,
-                               sg_logger=self.sg_logger,
-                               train_loader=self.train_loader,
-                               context_methods=self._get_context_methods(Phase.TRAIN_BATCH_END),
-                               ddp_silent_mode=self.ddp_silent_mode)
-
-        for batch_idx, batch_items in enumerate(progress_bar_train_loader):
-            batch_items = core_utils.tensor_container_to_device(batch_items, self.device, non_blocking=True)
-            inputs, targets, additional_batch_items = sg_model_utils.unpack_batch_items(batch_items)
-
-            if self.pre_prediction_callback is not None:
-                inputs, targets = self.pre_prediction_callback(inputs, targets, batch_idx)
-            # AUTOCAST IS ENABLED ONLY IF self.training_params.mixed_precision - IF enabled=False AUTOCAST HAS NO EFFECT
-            with autocast(enabled=self.training_params.mixed_precision):
-                # FORWARD PASS TO GET NETWORK'S PREDICTIONS
-                outputs = self.net(inputs)
-
-                # COMPUTE THE LOSS FOR BACK PROP + EXTRA METRICS COMPUTED DURING THE LOSS FORWARD PASS
-                loss, loss_log_items = self._get_losses(outputs, targets)
-
-            context.update_context(batch_idx=batch_idx,
-                                   inputs=inputs,
-                                   preds=outputs,
-                                   target=targets,
-                                   loss_log_items=loss_log_items,
-                                   **additional_batch_items)
-
-            self.phase_callback_handler(Phase.TRAIN_BATCH_END, context)
-
-            # LOG LR THAT WILL BE USED IN CURRENT EPOCH AND AFTER FIRST WARMUP/LR_SCHEDULER UPDATE BEFORE WEIGHT UPDATE
-            if not self.ddp_silent_mode and batch_idx == 0:
-                self._write_lrs(epoch)
-
-            self._backward_step(loss, epoch, batch_idx, context)
-
-            # COMPUTE THE RUNNING USER METRICS AND LOSS RUNNING ITEMS. RESULT TUPLE IS THEIR CONCATENATION.
-            logging_values = loss_avg_meter.average + get_metrics_results_tuple(self.train_metrics)
-            gpu_memory_utilization = torch.cuda.memory_cached() / 1E9 if torch.cuda.is_available() else 0
-
-            # RENDER METRICS PROGRESS
-            pbar_message_dict = get_train_loop_description_dict(logging_values,
-                                                                self.train_metrics,
-                                                                self.loss_logging_items_names,
-                                                                gpu_mem=gpu_memory_utilization)
-
-            progress_bar_train_loader.set_postfix(**pbar_message_dict)
-
-            # TODO: ITERATE BY MAX ITERS
-            # FOR INFINITE SAMPLERS WE MUST BREAK WHEN REACHING LEN ITERATIONS.
-            if self._infinite_train_loader and batch_idx == len(self.train_loader) - 1:
-                break
-
-        if not self.ddp_silent_mode:
-            self.sg_logger.upload()
-
-        self.train_monitored_values = sg_model_utils.update_monitored_values_dict(
-            monitored_values_dict=self.train_monitored_values, new_values_dict=pbar_message_dict)
-
-        return logging_values
-
-    def _get_losses(self, outputs: torch.Tensor, targets: torch.Tensor) -> Tuple[torch.Tensor, tuple]:
-        # GET THE OUTPUT OF THE LOSS FUNCTION
-        loss = self.criterion(outputs, targets)
-        if isinstance(loss, tuple):
-            loss, loss_logging_items = loss
-            # IF ITS NOT A TUPLE THE LOGGING ITEMS CONTAIN ONLY THE LOSS FOR BACKPROP (USER DEFINED LOSS RETURNS SCALAR)
-        else:
-            loss_logging_items = loss.unsqueeze(0).detach()
-
-        if len(loss_logging_items) != len(self.loss_logging_items_names):
-            raise ValueError("Loss output length must match loss_logging_items_names. Got " + str(
-                len(loss_logging_items)) + ', and ' + str(len(self.loss_logging_items_names)))
-        # RETURN AND THE LOSS LOGGING ITEMS COMPUTED DURING LOSS FORWARD PASS
-        return loss, loss_logging_items
-
-    def _backward_step(self, loss: torch.Tensor, epoch: int, batch_idx: int, context: PhaseContext, *args, **kwargs):
-        """
-        Run backprop on the loss and perform a step
-        :param loss: The value computed by the loss function
-        :param optimizer: An object that can perform a gradient step and zeroize model gradient
-        :param epoch: number of epoch the training is on
-        :param batch_idx: number of iteration inside the current epoch
-        :param context: current phase context
-        :return:
-        """
-        # SCALER IS ENABLED ONLY IF self.training_params.mixed_precision=True
-        self.scaler.scale(loss).backward()
-
-        # APPLY GRADIENT CLIPPING IF REQUIRED
-        if self.training_params.clip_grad_norm:
-            torch.nn.utils.clip_grad_norm_(self.net.parameters(), self.training_params.clip_grad_norm)
-
-        # ACCUMULATE GRADIENT FOR X BATCHES BEFORE OPTIMIZING
-        integrated_batches_num = batch_idx + len(self.train_loader) * epoch + 1
-
-        if integrated_batches_num % self.batch_accumulate == 0:
-            # SCALER IS ENABLED ONLY IF self.training_params.mixed_precision=True
-            self.scaler.step(self.optimizer)
-            self.scaler.update()
-
-            self.optimizer.zero_grad()
-            if self.ema:
-                self.ema_model.update(self.net, integrated_batches_num / (len(self.train_loader) * self.max_epochs))
-
-            # RUN PHASE CALLBACKS
-            self.phase_callback_handler(Phase.TRAIN_BATCH_STEP, context)
-
-    def _save_checkpoint(self, optimizer=None, epoch: int = None, validation_results_tuple: tuple = None,
-                         context: PhaseContext = None):
-        """
-        Save the current state dict as latest (always), best (if metric was improved), epoch# (if determined in training
-        params)
-        """
-        # WHEN THE validation_results_tuple IS NONE WE SIMPLY SAVE THE state_dict AS LATEST AND Return
-        if validation_results_tuple is None:
-            self.sg_logger.add_checkpoint(tag='ckpt_latest_weights_only.pth', state_dict={'net': self.net.state_dict()},
-                                          global_step=epoch)
-            return
-
-        # COMPUTE THE CURRENT metric
-        # IF idx IS A LIST - SUM ALL THE VALUES STORED IN THE LIST'S INDICES
-        metric = validation_results_tuple[self.metric_idx_in_results_tuple] if isinstance(
-            self.metric_idx_in_results_tuple, int) else \
-            sum([validation_results_tuple[idx] for idx in self.metric_idx_in_results_tuple])
-
-        # BUILD THE state_dict
-        state = {'net': self.net.state_dict(), 'acc': metric, 'epoch': epoch}
-        if optimizer is not None:
-            state['optimizer_state_dict'] = optimizer.state_dict()
-
-        if self.scaler is not None:
-            state['scaler_state_dict'] = self.scaler.state_dict()
-
-        if self.ema:
-            state['ema_net'] = self.ema_model.ema.state_dict()
-        # SAVES CURRENT MODEL AS ckpt_latest
-        self.sg_logger.add_checkpoint(tag='ckpt_latest.pth', state_dict=state, global_step=epoch)
-
-        # SAVE MODEL AT SPECIFIC EPOCHS DETERMINED BY save_ckpt_epoch_list
-        if epoch in self.training_params.save_ckpt_epoch_list:
-            self.sg_logger.add_checkpoint(tag=f'ckpt_epoch_{epoch}.pth', state_dict=state, global_step=epoch)
-
-        # OVERRIDE THE BEST CHECKPOINT AND best_metric IF metric GOT BETTER THAN THE PREVIOUS BEST
-        if (metric > self.best_metric and self.greater_metric_to_watch_is_better) or (
-                metric < self.best_metric and not self.greater_metric_to_watch_is_better):
-            # STORE THE CURRENT metric AS BEST
-            self.best_metric = metric
-            self._save_best_checkpoint(epoch, state)
-
-            # RUN PHASE CALLBACKS
-            self.phase_callback_handler(Phase.VALIDATION_END_BEST_EPOCH, context)
-
-            if isinstance(metric, torch.Tensor):
-                metric = metric.item()
-            logger.info("Best checkpoint overriden: validation " + self.metric_to_watch + ": " + str(metric))
-
-        if self.training_params.average_best_models:
-            net_for_averaging = self.ema_model.ema if self.ema else self.net
-            averaged_model_sd = self.model_weight_averaging.get_average_model(net_for_averaging,
-                                                                              validation_results_tuple=validation_results_tuple)
-            self.sg_logger.add_checkpoint(tag=self.average_model_checkpoint_filename,
-                                          state_dict={'net': averaged_model_sd}, global_step=epoch)
-
-    def _save_best_checkpoint(self, epoch, state):
-        self.sg_logger.add_checkpoint(tag=self.ckpt_best_name, state_dict=state, global_step=epoch)
-
-    # FIXME - we need to resolve flake8's 'function is too complex' for this function
-    def train(self, training_params: dict = dict()):  # noqa: C901
-        """
-
-        train - Trains the Model
-
-        IMPORTANT NOTE: Additional batch parameters can be added as a third item (optional) if a tuple is returned by
-          the data loaders, as dictionary. The phase context will hold the additional items, under an attribute with
-          the same name as the key in this dictionary. Then such items can be accessed through phase callbacks.
-
-
-            :param training_params:
-                - `max_epochs` : int
-
-                    Number of epochs to run training.
-
-                - `lr_updates` : list(int)
-
-                    List of fixed epoch numbers to perform learning rate updates when `lr_mode='step'`.
-
-                - `lr_decay_factor` : float
-
-                    Decay factor to apply to the learning rate at each update when `lr_mode='step'`.
-
-
-                -  `lr_mode` : str
-
-                    Learning rate scheduling policy, one of ['step','poly','cosine','function']. 'step' refers to
-                    constant updates at epoch numbers passed through `lr_updates`. 'cosine' refers to Cosine Anealing
-                    policy as mentioned in https://arxiv.org/abs/1608.03983. 'poly' refers to polynomial decrease i.e
-                    in each epoch iteration `self.lr = self.initial_lr * pow((1.0 - (current_iter / max_iter)),
-                    0.9)` 'function' refers to user defined learning rate scheduling function, that is passed through
-                    `lr_schedule_function`.
-
-                - `lr_schedule_function` : Union[callable,None]
-
-                    Learning rate scheduling function to be used when `lr_mode` is 'function'.
-
-                - `lr_warmup_epochs` : int (default=0)
-
-                    Number of epochs for learning rate warm up - see https://arxiv.org/pdf/1706.02677.pdf (Section 2.2).
-
-                - `cosine_final_lr_ratio` : float (default=0.01)
-                    Final learning rate ratio (only relevant when `lr_mode`='cosine'). The cosine starts from initial_lr and reaches
-                     initial_lr * cosine_final_lr_ratio in last epoch
-
-                - `inital_lr` : float
-
-                    Initial learning rate.
-
-                - `loss` : Union[nn.module, str]
-
-                    Loss function for training.
-                    One of SuperGradient's built in options:
-
-                              "cross_entropy": LabelSmoothingCrossEntropyLoss,
-                              "mse": MSELoss,
-                              "r_squared_loss": RSquaredLoss,
-                              "detection_loss": YoLoV3DetectionLoss,
-                              "shelfnet_ohem_loss": ShelfNetOHEMLoss,
-                              "shelfnet_se_loss": ShelfNetSemanticEncodingLoss,
-                              "ssd_loss": SSDLoss,
-
-
-                    or user defined nn.module loss function.
-
-                    IMPORTANT: forward(...) should return a (loss, loss_items) tuple where loss is the tensor used
-                    for backprop (i.e what your original loss function returns), and loss_items should be a tensor of
-                    shape (n_items), of values computed during the forward pass which we desire to log over the
-                    entire epoch. For example- the loss itself should always be logged. Another example is a scenario
-                    where the computed loss is the sum of a few components we would like to log- these entries in
-                    loss_items).
-
-                    When training, set the loss_logging_items_names parameter in train_params to be a list of
-                    strings, of length n_items who's ith element is the name of the ith entry in loss_items. Then
-                    each item will be logged, rendered on tensorboard and "watched" (i.e saving model checkpoints
-                    according to it).
-
-                    Since running logs will save the loss_items in some internal state, it is recommended that
-                    loss_items are detached from their computational graph for memory efficiency.
-
-                - `optimizer` : Union[str, torch.optim.Optimizer]
-
-                    Optimization algorithm. One of ['Adam','SGD','RMSProp'] corresponding to the torch.optim
-                    optimzers implementations, or any object that implements torch.optim.Optimizer.
-
-                - `criterion_params` : dict
-
-                    Loss function parameters.
-
-                - `optimizer_params` : dict
-                    When `optimizer` is one of ['Adam','SGD','RMSProp'], it will be initialized with optimizer_params.
-
-                    (see https://pytorch.org/docs/stable/optim.html for the full list of
-                    parameters for each optimizer).
-
-                - `train_metrics_list` : list(torchmetrics.Metric)
-
-                    Metrics to log during training. For more information on torchmetrics see
-                    https://torchmetrics.rtfd.io/en/latest/.
-
-
-                - `valid_metrics_list` : list(torchmetrics.Metric)
-
-                    Metrics to log during validation/testing. For more information on torchmetrics see
-                    https://torchmetrics.rtfd.io/en/latest/.
-
-
-                - `loss_logging_items_names` : list(str)
-
-                    The list of names/titles for the outputs returned from the loss functions forward pass (reminder-
-                    the loss function should return the tuple (loss, loss_items)). These names will be used for
-                    logging their values.
-
-                - `metric_to_watch` : str (default="Accuracy")
-
-                    will be the metric which the model checkpoint will be saved according to, and can be set to any
-                    of the following:
-
-                        a metric name (str) of one of the metric objects from the valid_metrics_list
-
-                        a "metric_name" if some metric in valid_metrics_list has an attribute component_names which
-                        is a list referring to the names of each entry in the output metric (torch tensor of size n)
-
-                        one of "loss_logging_items_names" i.e which will correspond to an item returned during the
-                        loss function's forward pass.
-
-                    At the end of each epoch, if a new best metric_to_watch value is achieved, the models checkpoint
-                    is saved in YOUR_PYTHON_PATH/checkpoints/ckpt_best.pth
-
-                - `greater_metric_to_watch_is_better` : bool
-
-                    When choosing a model's checkpoint to be saved, the best achieved model is the one that maximizes the
-                     metric_to_watch when this parameter is set to True, and a one that minimizes it otherwise.
-
-                - `ema` : bool (default=False)
-
-                    Whether to use Model Exponential Moving Average (see
-                    https://github.com/rwightman/pytorch-image-models ema implementation)
-
-                - `batch_accumulate` : int (default=1)
-
-                    Number of batches to accumulate before every backward pass.
-
-                - `ema_params` : dict
-
-                    Parameters for the ema model.
-
-                - `zero_weight_decay_on_bias_and_bn` : bool (default=False)
-
-                    Whether to apply weight decay on batch normalization parameters or not (ignored when the passed
-                    optimizer has already been initialized).
-
-
-                - `load_opt_params` : bool (default=True)
-
-                    Whether to load the optimizers parameters as well when loading a model's checkpoint.
-
-                - `run_validation_freq` : int (default=1)
-
-                    The frequency in which validation is performed during training (i.e the validation is ran every
-                     `run_validation_freq` epochs.
-
-                - `save_model` : bool (default=True)
-
-                    Whether to save the model checkpoints.
-
-                - `silent_mode` : bool
-
-                    Silents the print outs.
-
-                - `mixed_precision` : bool
-
-                    Whether to use mixed precision or not.
-
-                - `save_ckpt_epoch_list` : list(int) (default=[])
-
-                    List of fixed epoch indices the user wishes to save checkpoints in.
-
-                - `average_best_models` : bool (default=False)
-
-                    If set, a snapshot dictionary file and the average model will be saved / updated at every epoch
-                    and evaluated only when training is completed. The snapshot file will only be deleted upon
-                    completing the training. The snapshot dict will be managed on cpu.
-
-                - `precise_bn` : bool (default=False)
-
-                    Whether to use precise_bn calculation during the training.
-
-                - `precise_bn_batch_size` : int (default=None)
-
-                    The effective batch size we want to calculate the batchnorm on. For example, if we are training a model
-                    on 8 gpus, with a batch of 128 on each gpu, a good rule of thumb would be to give it 8192
-                    (ie: effective_batch_size * num_gpus = batch_per_gpu * num_gpus * num_gpus).
-                    If precise_bn_batch_size is not provided in the training_params, the latter heuristic will be taken.
-
-                - `seed` : int (default=42)
-
-                    Random seed to be set for torch, numpy, and random. When using DDP each process will have it's seed
-                    set to seed + rank.
-
-
-                - `log_installed_packages` : bool (default=False)
-
-                    When set, the list of all installed packages (and their versions) will be written to the tensorboard
-                     and logfile (useful when trying to reproduce results).
-
-                - `dataset_statistics` : bool (default=False)
-
-                    Enable a statistic analysis of the dataset. If set to True the dataset will be analyzed and a report
-                    will be added to the tensorboard along with some sample images from the dataset. Currently only
-                    detection datasets are supported for analysis.
-
-                -  `save_full_train_log` : bool (default=False)
-
-                    When set, a full log (of all super_gradients modules, including uncaught exceptions from any other
-                     module) of the training will be saved in the checkpoint directory under full_train_log.log
-
-                -  `sg_logger` : Union[AbstractSGLogger, str] (defauls=base_sg_logger)
-
-                    Define the SGLogger object for this training process. The SGLogger handles all disk writes, logs, TensorBoard, remote logging
-                    and remote storage. By overriding the default base_sg_logger, you can change the storage location, support external monitoring and logging
-                    or support remote storage.
-
-                -   `sg_logger_params` : dict
-
-                    SGLogger parameters
-
-                -   `clip_grad_norm` : float
-
-                    Defines a maximal L2 norm of the gradients. Values which exceed the given value will be clipped
-
-                -   `lr_cooldown_epochs` : int (default=0)
-
-                    Number of epochs to cooldown LR (i.e the last epoch from scheduling view point=max_epochs-cooldown).
-
-                -   `pre_prediction_callback` : Callable (default=None)
-
-                     When not None, this callback will be applied to images and targets, and returning them to be used
-                      for the forward pass, and further computations. Args for this callable should be in the order
-                      (inputs, targets, batch_idx) returning modified_inputs, modified_targets
-
-                -   `ckpt_best_name` : str (default='ckpt_best.pth')
-
-                    The best checkpoint (according to metric_to_watch) will be saved under this filename in the checkpoints directory.
-
-                -   `enable_qat`: bool (default=False)
-
-                    Adds a QATCallback to the phase callbacks, that triggers quantization aware training starting from
-                     qat_params["start_epoch"]
-
-                -   `qat_params`: dict-like object with the following key/values:
-
-                        start_epoch: int, first epoch to start QAT.
-
-                        quant_modules_calib_method: str, One of [percentile, mse, entropy, max]. Statistics method for amax
-                         computation of the quantized modules (default=percentile).
-
-                        per_channel_quant_modules: bool, whether quant modules should be per channel (default=False).
-
-                        calibrate: bool, whether to perfrom calibration (default=False).
-
-                        calibrated_model_path: str, path to a calibrated checkpoint (default=None).
-
-                        calib_data_loader: torch.utils.data.DataLoader, data loader of the calibration dataset. When None,
-                         context.train_loader will be used (default=None).
-
-                        num_calib_batches: int, number of batches to collect the statistics from.
-
-                        percentile: float, percentile value to use when SgModel,quant_modules_calib_method='percentile'.
-                         Discarded when other methods are used (Default=99.99).
-
-
-        :return:
-        """
-        global logger
-
-        if self.net is None:
-            raise Exception('Model', 'No model found')
-        if self.dataset_interface is None and self.train_loader is None:
-            raise Exception('Data', 'No dataset found')
-
-        self.training_params = TrainingParams()
-        self.training_params.override(**training_params)
-
-        # SET RANDOM SEED
-        random_seed(is_ddp=self.multi_gpu == MultiGPUMode.DISTRIBUTED_DATA_PARALLEL,
-                    device=self.device, seed=self.training_params.seed)
-
-        silent_mode = self.training_params.silent_mode or self.ddp_silent_mode
-        # METRICS
-        self._set_train_metrics(train_metrics_list=self.training_params.train_metrics_list)
-        self._set_valid_metrics(valid_metrics_list=self.training_params.valid_metrics_list)
-        self.loss_logging_items_names = self.training_params.loss_logging_items_names
-
-        self.results_titles = ["Train_" + t for t in
-                               self.loss_logging_items_names + get_metrics_titles(self.train_metrics)] + \
-                              ["Valid_" + t for t in
-                               self.loss_logging_items_names + get_metrics_titles(self.valid_metrics)]
-
-        # Store the metric to follow (loss\accuracy) and initialize as the worst value
-        self.metric_to_watch = self.training_params.metric_to_watch
-        self.greater_metric_to_watch_is_better = self.training_params.greater_metric_to_watch_is_better
-        self.metric_idx_in_results_tuple = (
-            self.loss_logging_items_names + get_metrics_titles(self.valid_metrics)).index(self.metric_to_watch)
-
-        # Instantiate the values to monitor (loss/metric)
-        for loss in self.loss_logging_items_names:
-            self.train_monitored_values[loss] = MonitoredValue(name=loss, greater_is_better=False)
-            self.valid_monitored_values[loss] = MonitoredValue(name=loss, greater_is_better=False)
-        self.valid_monitored_values[self.metric_to_watch] = MonitoredValue(name=self.metric_to_watch,
-                                                                           greater_is_better=True)
-
-        # Allowing loading instantiated loss or string
-        if isinstance(self.training_params.loss, str):
-            criterion_cls = LOSSES[self.training_params.loss]
-            self.criterion = criterion_cls(**self.training_params.criterion_params)
-
-        elif isinstance(self.training_params.loss, Mapping):
-            self.criterion = LossesFactory().get(self.training_params.loss)
-
-        elif isinstance(self.training_params.loss, nn.Module):
-            self.criterion = self.training_params.loss
-
-        self.criterion.to(self.device)
-
-        self.max_epochs = self.training_params.max_epochs
-
-        self.ema = self.training_params.ema
-
-        self.precise_bn = self.training_params.precise_bn
-        self.precise_bn_batch_size = self.training_params.precise_bn_batch_size
-
-        self.batch_accumulate = self.training_params.batch_accumulate
-        num_batches = len(self.train_loader)
-
-        if self.ema:
-            ema_params = self.training_params.ema_params
-            logger.info(f'Using EMA with params {ema_params}')
-            self.ema_model = self._instantiate_ema_model(**ema_params)
-            self.ema_model.updates = self.start_epoch * num_batches // self.batch_accumulate
-            if self.load_checkpoint:
-                if 'ema_net' in self.checkpoint.keys():
-                    self.ema_model.ema.load_state_dict(self.checkpoint['ema_net'])
-                else:
-                    self.ema = False
-                    logger.warning(
-                        "[Warning] Checkpoint does not include EMA weights, continuing training without EMA.")
-
-        self.run_validation_freq = self.training_params.run_validation_freq
-        validation_results_tuple = (0, 0)
-        inf_time = 0
-        timer = core_utils.Timer(self.device)
-
-        # IF THE LR MODE IS NOT DEFAULT TAKE IT FROM THE TRAINING PARAMS
-        self.lr_mode = self.training_params.lr_mode
-        load_opt_params = self.training_params.load_opt_params
-
-        self.phase_callbacks = self.training_params.phase_callbacks or []
-        self.phase_callbacks = ListFactory(CallbacksFactory()).get(self.phase_callbacks)
-
-        if self.lr_mode is not None:
-            sg_lr_callback_cls = LR_SCHEDULERS_CLS_DICT[self.lr_mode]
-            self.phase_callbacks.append(sg_lr_callback_cls(train_loader_len=len(self.train_loader),
-                                                           net=self.net,
-                                                           training_params=self.training_params,
-                                                           update_param_groups=self.update_param_groups,
-                                                           **self.training_params.to_dict()))
-        if self.training_params.lr_warmup_epochs > 0:
-            warmup_mode = self.training_params.warmup_mode
-            if isinstance(warmup_mode, str):
-                warmup_callback_cls = LR_WARMUP_CLS_DICT[warmup_mode]
-            elif isinstance(warmup_mode, type) and issubclass(warmup_mode, LRCallbackBase):
-                warmup_callback_cls = warmup_mode
-            else:
-                raise RuntimeError('warmup_mode has to be either a name of a mode (str) or a subclass of PhaseCallback')
-            self.phase_callbacks.append(warmup_callback_cls(train_loader_len=len(self.train_loader),
-                                                            net=self.net,
-                                                            training_params=self.training_params,
-                                                            update_param_groups=self.update_param_groups,
-                                                            **self.training_params.to_dict()))
-
-        self._add_metrics_update_callback(Phase.TRAIN_BATCH_END)
-        self._add_metrics_update_callback(Phase.VALIDATION_BATCH_END)
-
-        # ADD CALLBACK FOR QAT
-        self.enable_qat = core_utils.get_param(self.training_params, "enable_qat", False)
-        if self.enable_qat:
-            self.qat_params = core_utils.get_param(self.training_params, "qat_params")
-            if self.qat_params is None:
-                raise ValueError("Must pass QAT params when enable_qat=True")
-            self.phase_callbacks.append(QATCallback(**self.qat_params))
-
-        self.phase_callback_handler = CallbackHandler(callbacks=self.phase_callbacks)
-
-        if not self.ddp_silent_mode:
-            self._initialize_sg_logger_objects()
-
-            if self.training_params.dataset_statistics:
-                dataset_statistics_logger = DatasetStatisticsTensorboardLogger(self.sg_logger)
-                dataset_statistics_logger.analyze(self.train_loader, all_classes=self.classes,
-                                                  title="Train-set", anchors=self.net.module.arch_params.anchors)
-                dataset_statistics_logger.analyze(self.valid_loader, all_classes=self.classes,
-                                                  title="val-set")
-            # AVERAGE BEST 10 MODELS PARAMS
-            if self.training_params.average_best_models:
-                self.model_weight_averaging = ModelWeightAveraging(self.checkpoints_dir_path,
-                                                                   greater_is_better=self.greater_metric_to_watch_is_better,
-                                                                   source_ckpt_folder_name=self.source_ckpt_folder_name,
-                                                                   metric_to_watch=self.metric_to_watch,
-                                                                   metric_idx=self.metric_idx_in_results_tuple,
-                                                                   load_checkpoint=self.load_checkpoint,
-                                                                   model_checkpoints_location=self.model_checkpoints_location)
-        if self.training_params.save_full_train_log and not self.ddp_silent_mode:
-            logger = get_logger(__name__,
-                                training_log_path=self.sg_logger.log_file_path.replace('.txt', 'full_train_log.log'))
-            sg_model_utils.log_uncaught_exceptions(logger)
-
-        if not self.load_checkpoint or self.load_weights_only:
-            # WHEN STARTING TRAINING FROM SCRATCH, DO NOT LOAD OPTIMIZER PARAMS (EVEN IF LOADING BACKBONE)
-            self.start_epoch = 0
-            self._reset_best_metric()
-            load_opt_params = False
-
-        if isinstance(self.training_params.optimizer, str) or \
-                (inspect.isclass(self.training_params.optimizer) and issubclass(self.training_params.optimizer, torch.optim.Optimizer)):
-            self.optimizer = build_optimizer(net=self.net, lr=self.training_params.initial_lr,
-                                             training_params=self.training_params)
-        elif isinstance(self.training_params.optimizer, torch.optim.Optimizer):
-            self.optimizer = self.training_params.optimizer
-        else:
-            raise UnsupportedOptimizerFormat()
-
-        # VERIFY GRADIENT CLIPPING VALUE
-        if self.training_params.clip_grad_norm is not None and self.training_params.clip_grad_norm <= 0:
-            raise TypeError('Params', 'Invalid clip_grad_norm')
-
-        if self.load_checkpoint and load_opt_params:
-            self.optimizer.load_state_dict(self.checkpoint['optimizer_state_dict'])
-
-        self.pre_prediction_callback = CallbacksFactory().get(self.training_params.pre_prediction_callback)
-
-        self._initialize_mixed_precision(self.training_params.mixed_precision)
-
-        self._infinite_train_loader = (hasattr(self.train_loader, "sampler") and isinstance(self.train_loader.sampler, InfiniteSampler)) or\
-                                      (hasattr(self.train_loader, "batch_sampler") and isinstance(self.train_loader.batch_sampler.sampler, InfiniteSampler))
-
-        self.ckpt_best_name = self.training_params.ckpt_best_name
-
-        context = PhaseContext(optimizer=self.optimizer,
-                               net=self.net,
-                               experiment_name=self.experiment_name,
-                               ckpt_dir=self.checkpoints_dir_path,
-                               criterion=self.criterion,
-                               lr_warmup_epochs=self.training_params.lr_warmup_epochs,
-                               sg_logger=self.sg_logger,
-                               train_loader=self.train_loader,
-                               valid_loader=self.valid_loader,
-                               training_params=self.training_params,
-                               ddp_silent_mode=self.ddp_silent_mode,
-                               checkpoint_params=self.checkpoint_params,
-                               architecture=self.architecture,
-                               arch_params=self.arch_params,
-                               metric_idx_in_results_tuple=self.metric_idx_in_results_tuple,
-                               metric_to_watch=self.metric_to_watch,
-                               device=self.device,
-                               context_methods=self._get_context_methods(Phase.PRE_TRAINING)
-                               )
-
-        self.phase_callback_handler(Phase.PRE_TRAINING, context)
-
-        try:
-            # HEADERS OF THE TRAINING PROGRESS
-            if not silent_mode:
-                logger.info(
-                    f'Started training for {self.max_epochs - self.start_epoch} epochs ({self.start_epoch}/'f'{self.max_epochs - 1})\n')
-            for epoch in range(self.start_epoch, self.max_epochs):
-                if context.stop_training:
-                    logger.info("Request to stop training has been received, stopping training")
-                    break
-
-                # Phase.TRAIN_EPOCH_START
-                # RUN PHASE CALLBACKS
-                context.update_context(epoch=epoch)
-                self.phase_callback_handler(Phase.TRAIN_EPOCH_START, context)
-
-                # IN DDP- SET_EPOCH WILL CAUSE EVERY PROCESS TO BE EXPOSED TO THE ENTIRE DATASET BY SHUFFLING WITH A
-                # DIFFERENT SEED EACH EPOCH START
-                if self.multi_gpu == MultiGPUMode.DISTRIBUTED_DATA_PARALLEL and hasattr(self.train_loader, "sampler")\
-                        and hasattr(self.train_loader.sampler, "set_epoch"):
-                    self.train_loader.sampler.set_epoch(epoch)
-
-                train_metrics_tuple = self._train_epoch(epoch=epoch, silent_mode=silent_mode)
-
-                # Phase.TRAIN_EPOCH_END

# PACKAGE IMPORTS FOR EXTERNAL USAGE

from super_gradients.training.sg_trainer.sg_trainer import Trainer, MultiGPUMode, StrictLoad

__all__ = ['Trainer', 'MultiGPUMode', 'StrictLoad']

import inspect
import os
import sys
from copy import deepcopy
from typing import Union, Tuple, Mapping, List, Any

import hydra
import numpy as np
import pkg_resources
import torch
from omegaconf import DictConfig
from torch import nn
from torch.utils.data import DataLoader, DistributedSampler
from torch.cuda.amp import GradScaler, autocast
from torchmetrics import MetricCollection
from tqdm import tqdm
from piptools.scripts.sync import _get_installed_distributions

from super_gradients.common.factories.callbacks_factory import CallbacksFactory
from super_gradients.common.data_types.enum import MultiGPUMode, StrictLoad, EvaluationType
from super_gradients.training.models.all_architectures import ARCHITECTURES
from super_gradients.common.decorators.factory_decorator import resolve_param
from super_gradients.common.environment import env_helpers
from super_gradients.common.abstractions.abstract_logger import get_logger
from super_gradients.common.factories.datasets_factory import DatasetsFactory
from super_gradients.common.factories.list_factory import ListFactory
from super_gradients.common.factories.losses_factory import LossesFactory
from super_gradients.common.factories.metrics_factory import MetricsFactory
from super_gradients.common.sg_loggers import SG_LOGGERS
from super_gradients.common.sg_loggers.abstract_sg_logger import AbstractSGLogger
from super_gradients.common.sg_loggers.base_sg_logger import BaseSGLogger
from super_gradients.training import utils as core_utils
from super_gradients.training.models import SgModule
from super_gradients.training.pretrained_models import PRETRAINED_NUM_CLASSES
from super_gradients.training.utils import sg_trainer_utils
from super_gradients.training.utils.quantization_utils import QATCallback
from super_gradients.training.utils.sg_trainer_utils import MonitoredValue, parse_args
from super_gradients.training.exceptions.sg_trainer_exceptions import UnsupportedOptimizerFormat, \
    IllegalDataloaderInitialization
from super_gradients.training.datasets import DatasetInterface
from super_gradients.training.losses import LOSSES
from super_gradients.training.metrics.metric_utils import get_metrics_titles, get_metrics_results_tuple, \
    get_logging_values, \
    get_metrics_dict, get_train_loop_description_dict
from super_gradients.training.params import TrainingParams
from super_gradients.training.utils.detection_utils import DetectionPostPredictionCallback
from super_gradients.training.utils.distributed_training_utils import MultiGPUModeAutocastWrapper, \
    reduce_results_tuple_for_ddp, compute_precise_bn_stats
from super_gradients.training.utils.ema import ModelEMA
from super_gradients.training.utils.optimizer_utils import build_optimizer
from super_gradients.training.utils.weight_averaging_utils import ModelWeightAveraging
from super_gradients.training.metrics import Accuracy, Top5
from super_gradients.training.utils import random_seed
from super_gradients.training.utils.checkpoint_utils import get_ckpt_local_path, read_ckpt_state_dict, \
    load_checkpoint_to_model, load_pretrained_weights
from super_gradients.training.datasets.datasets_utils import DatasetStatisticsTensorboardLogger
from super_gradients.training.utils.callbacks import CallbackHandler, Phase, LR_SCHEDULERS_CLS_DICT, PhaseContext, \
    MetricsUpdateCallback, LR_WARMUP_CLS_DICT, ContextSgMethods, LRCallbackBase
from super_gradients.common.environment import environment_config
from super_gradients.training.utils import HpmStruct
from super_gradients.training.datasets.samplers.infinite_sampler import InfiniteSampler

logger = get_logger(__name__)


class Trainer:
    """
    SuperGradient Model - Base Class for Sg Models

    Methods
    -------
    train(max_epochs : int, initial_epoch : int, save_model : bool)
        the main function used for the training, h.p. updating, logging etc.

    predict(idx : int)
        returns the predictions and label of the current inputs

    test(epoch : int, idx : int, save : bool):
        returns the test loss, accuracy and runtime
    """

    def __init__(self, experiment_name: str, device: str = None, multi_gpu: Union[MultiGPUMode, str] = MultiGPUMode.OFF,
                 model_checkpoints_location: str = 'local',
                 overwrite_local_checkpoint: bool = True, ckpt_name: str = 'ckpt_latest.pth',
                 post_prediction_callback: DetectionPostPredictionCallback = None, ckpt_root_dir: str = None,
                 train_loader: DataLoader = None, valid_loader: DataLoader = None, test_loader: DataLoader = None,
                 classes: List[Any] = None):
        """

        :param experiment_name:                      Used for logging and loading purposes
        :param device:                          If equal to 'cpu' runs on the CPU otherwise on GPU
        :param multi_gpu:                       If True, runs on all available devices
        :param model_checkpoints_location:      If set to 's3' saves the Checkpoints in AWS S3
                                                otherwise saves the Checkpoints Locally
        :param overwrite_local_checkpoint:      If set to False keeps the current local checkpoint when importing
                                                checkpoint from cloud service, otherwise overwrites the local checkpoints file
        :param ckpt_name:                       The Checkpoint to Load
        :param ckpt_root_dir:                   Local root directory path where all experiment logging directories will
                                                reside. When none is give, it is assumed that
                                                pkg_resources.resource_filename('checkpoints', "") exists and will be used.
        :param train_loader:                    Training set Dataloader instead of using DatasetInterface, must pass "valid_loader"
                                                and "classes" along with it
        :param valid_loader:                    Validation set Dataloader
        :param test_loader:                     Test set Dataloader
        :param classes:                         List of class labels

        """
        # SET THE EMPTY PROPERTIES
        self.net, self.architecture, self.arch_params, self.dataset_interface = None, None, None, None
        self.device, self.multi_gpu = None, None
        self.ema = None
        self.ema_model = None
        self.sg_logger = None
        self.update_param_groups = None
        self.post_prediction_callback = None
        self.criterion = None
        self.training_params = None
        self.scaler = None
        self.phase_callbacks = None
        self.checkpoint_params = None
        self.pre_prediction_callback = None

        # SET THE DEFAULT PROPERTIES
        self.half_precision = False
        self.load_checkpoint = False
        self.load_backbone = False
        self.load_weights_only = False
        self.ddp_silent_mode = False
        self.source_ckpt_folder_name = None
        self.model_weight_averaging = None
        self.average_model_checkpoint_filename = 'average_model.pth'
        self.start_epoch = 0
        self.best_metric = np.inf
        self.external_checkpoint_path = None
        self.strict_load = StrictLoad.ON
        self.load_ema_as_net = False
        self.ckpt_best_name = 'ckpt_best.pth'
        self.enable_qat = False
        self.qat_params = {}
        self._infinite_train_loader = False

        # DETERMINE THE LOCATION OF THE LOSS AND ACCURACY IN THE RESULTS TUPLE OUTPUTED BY THE TEST
        self.loss_idx_in_results_tuple, self.acc_idx_in_results_tuple = None, None

        # METRICS
        self.loss_logging_items_names = None
        self.train_metrics = None
        self.valid_metrics = None
        self.greater_metric_to_watch_is_better = None

        # SETTING THE PROPERTIES FROM THE CONSTRUCTOR
        self.experiment_name = experiment_name
        self.ckpt_name = ckpt_name
        self.overwrite_local_checkpoint = overwrite_local_checkpoint
        self.model_checkpoints_location = model_checkpoints_location
        self._set_dataset_properties(classes, test_loader, train_loader, valid_loader)

        # CREATING THE LOGGING DIR BASED ON THE INPUT PARAMS TO PREVENT OVERWRITE OF LOCAL VERSION
        if ckpt_root_dir:
            self.checkpoints_dir_path = os.path.join(ckpt_root_dir, self.experiment_name)
        elif pkg_resources.resource_exists("checkpoints", ""):
            self.checkpoints_dir_path = pkg_resources.resource_filename('checkpoints', self.experiment_name)
        else:
            raise ValueError("Illegal checkpoints directory: pass ckpt_root_dir that exists, or add 'checkpoints' to"
                             "resources.")

        # INITIALIZE THE DEVICE FOR THE MODEL
        self._initialize_device(requested_device=device, requested_multi_gpu=multi_gpu)

        self.post_prediction_callback = post_prediction_callback
        # SET THE DEFAULTS
        # TODO: SET DEFAULT TRAINING PARAMS FOR EACH TASK

        default_results_titles = ['Train Loss', 'Train Acc', 'Train Top5', 'Valid Loss', 'Valid Acc', 'Valid Top5']

        self.results_titles = default_results_titles

        self.loss_idx_in_results_tuple, self.acc_idx_in_results_tuple = 0, 1
        default_train_metrics, default_valid_metrics = MetricCollection([Accuracy(), Top5()]), MetricCollection(
            [Accuracy(), Top5()])

        default_loss_logging_items_names = ["Loss"]

        self.train_metrics, self.valid_metrics = default_train_metrics, default_valid_metrics
        self.loss_logging_items_names = default_loss_logging_items_names

        self.train_monitored_values = {}
        self.valid_monitored_values = {}

    @classmethod
    def train_from_config(cls, cfg: Union[DictConfig, dict]) -> None:
        """
        Trains according to cfg recipe configuration.

        @param cfg: The parsed DictConfig from yaml recipe files or a dictionary
        @return: output of trainer.train(...) (i.e results tuple)
        """
        # INSTANTIATE ALL OBJECTS IN CFG
        cfg = hydra.utils.instantiate(cfg)

        kwargs = parse_args(cfg, cls.__init__)

        trainer = Trainer(**kwargs)

        # CONNECT THE DATASET INTERFACE WITH DECI MODEL
        trainer.connect_dataset_interface(cfg.dataset_interface, data_loader_num_workers=cfg.data_loader_num_workers)

        # BUILD NETWORK
        trainer.build_model(cfg.architecture, arch_params=cfg.arch_params, checkpoint_params=cfg.checkpoint_params)

        # TRAIN
        trainer.train(training_params=cfg.training_hyperparams)

    def _set_dataset_properties(self, classes, test_loader, train_loader, valid_loader):
        if any([train_loader, valid_loader, classes]) and not all([train_loader, valid_loader, classes]):
            raise IllegalDataloaderInitialization()

        dataset_params = {"batch_size": train_loader.batch_size if train_loader else None,
                          "val_batch_size": valid_loader.batch_size if valid_loader else None,
                          "test_batch_size": test_loader.batch_size if test_loader else None,
                          "dataset_dir": None,
                          "s3_link": None}

        if train_loader and self.multi_gpu == MultiGPUMode.DISTRIBUTED_DATA_PARALLEL:
            if not all([isinstance(train_loader.sampler, DistributedSampler),
                        isinstance(valid_loader.sampler, DistributedSampler),
                        test_loader is None or isinstance(test_loader.sampler, DistributedSampler)]):
                logger.warning("DDP training was selected but the dataloader samplers are not of type DistributedSamplers")

        self.dataset_params, self.train_loader, self.valid_loader, self.test_loader, self.classes = \
            HpmStruct(**dataset_params), train_loader, valid_loader, test_loader, classes

    @resolve_param('dataset_interface', DatasetsFactory())
    def connect_dataset_interface(self, dataset_interface: DatasetInterface, data_loader_num_workers: int = 8):
        """
        :param dataset_interface: DatasetInterface object
        :param data_loader_num_workers: The number of threads to initialize the Data Loaders with
            The dataset to be connected
        """
        if self.train_loader:
            logger.warning("Overriding the dataloaders that Trainer was initialized with")
        self.dataset_interface = dataset_interface
        self.train_loader, self.valid_loader, self.test_loader, self.classes = \
            self.dataset_interface.get_data_loaders(batch_size_factor=self.num_devices,
                                                    num_workers=data_loader_num_workers,
                                                    distributed_sampler=self.multi_gpu == MultiGPUMode.DISTRIBUTED_DATA_PARALLEL)

        self.dataset_params = self.dataset_interface.get_dataset_params()

    # FIXME - we need to resolve flake8's 'function is too complex' for this function
    def build_model(self,  # noqa: C901 - too complex
                    architecture: Union[str, nn.Module],
                    arch_params={}, checkpoint_params={}, *args, **kwargs):
        """
        :param architecture:               Defines the network's architecture from models/ALL_ARCHITECTURES
        :param arch_params:                Architecture H.P. e.g.: block, num_blocks, num_classes, etc.
        :param checkpoint_params:          Dictionary like object with the following key:values:

            load_checkpoint:            Load a pre-trained checkpoint
            strict_load:                See StrictLoad class documentation for details.
            source_ckpt_folder_name:    folder name to load the checkpoint from (self.experiment_name if none is given)
            load_weights_only:          loads only the weight from the checkpoint and zeroize the training params
            load_backbone:              loads the provided checkpoint to self.net.backbone instead of self.net
            external_checkpoint_path:   The path to the external checkpoint to be loaded. Can be absolute or relative
                                               (ie: path/to/checkpoint.pth). If provided, will automatically attempt to
                                               load the checkpoint even if the load_checkpoint flag is not provided.

        """
        if 'num_classes' not in arch_params.keys():
            if self.classes is None and self.dataset_interface is None:
                raise Exception('Error', 'Number of classes not defined in arch params and dataset is not defined')
            else:
                arch_params['num_classes'] = len(self.classes)

        self.arch_params = core_utils.HpmStruct(**arch_params)
        self.checkpoint_params = core_utils.HpmStruct(**checkpoint_params)

        self.net = self._instantiate_net(architecture, self.arch_params, checkpoint_params, *args, **kwargs)

        # SAVE THE ARCHITECTURE FOR NEURAL ARCHITECTURE SEARCH

        self.architecture = architecture

        self._net_to_device()

        # SET THE FLAG FOR DIFFERENT PARAMETER GROUP OPTIMIZER UPDATE
        self.update_param_groups = hasattr(self.net.module, 'update_param_groups')

        self._load_checkpoint_to_model()

    def _set_ckpt_loading_attributes(self):
        """
        Sets checkpoint loading related attributes according to self.checkpoint_params
        """
        self.checkpoint = {}
        self.strict_load = core_utils.get_param(self.checkpoint_params, 'strict_load', default_val=StrictLoad.ON)
        self.load_ema_as_net = core_utils.get_param(self.checkpoint_params, 'load_ema_as_net', default_val=False)
        self.source_ckpt_folder_name = core_utils.get_param(self.checkpoint_params, 'source_ckpt_folder_name')
        self.load_checkpoint = core_utils.get_param(self.checkpoint_params, 'load_checkpoint', default_val=False)
        self.load_backbone = core_utils.get_param(self.checkpoint_params, 'load_backbone', default_val=False)
        self.external_checkpoint_path = core_utils.get_param(self.checkpoint_params, 'external_checkpoint_path')
        if self.load_checkpoint or self.external_checkpoint_path:
            self.load_weights_only = core_utils.get_param(self.checkpoint_params, 'load_weights_only',
                                                          default_val=False)
        self.ckpt_name = core_utils.get_param(self.checkpoint_params, 'ckpt_name', default_val=self.ckpt_name)

    def _net_to_device(self):
        """
        Manipulates self.net according to self.multi_gpu
        """
        self.net.to(self.device)

        # FOR MULTI-GPU TRAINING (not distributed)
        self.arch_params.sync_bn = core_utils.get_param(self.arch_params, 'sync_bn', default_val=False)
        if self.multi_gpu == MultiGPUMode.DATA_PARALLEL:
            self.net = torch.nn.DataParallel(self.net, device_ids=self.device_ids)
        elif self.multi_gpu == MultiGPUMode.DISTRIBUTED_DATA_PARALLEL:
            if self.arch_params.sync_bn:
                if not self.ddp_silent_mode:
                    logger.info('DDP - Using Sync Batch Norm... Training time will be affected accordingly')
                self.net = torch.nn.SyncBatchNorm.convert_sync_batchnorm(self.net).to(self.device)

            local_rank = int(self.device.split(':')[1])
            self.net = torch.nn.parallel.DistributedDataParallel(self.net,
                                                                 device_ids=[local_rank],
                                                                 output_device=local_rank,
                                                                 find_unused_parameters=True)

        else:
            self.net = core_utils.WrappedModel(self.net)

    def _train_epoch(self, epoch: int, silent_mode: bool = False) -> tuple:
        """
        train_epoch - A single epoch training procedure
            :param optimizer:   The optimizer for the network
            :param epoch:       The current epoch
            :param silent_mode: No verbosity
        """
        # SET THE MODEL IN training STATE
        self.net.train()
        # THE DISABLE FLAG CONTROLS WHETHER THE PROGRESS BAR IS SILENT OR PRINTS THE LOGS
        progress_bar_train_loader = tqdm(self.train_loader, bar_format="{l_bar}{bar:10}{r_bar}", dynamic_ncols=True,
                                         disable=silent_mode)
        progress_bar_train_loader.set_description(f"Train epoch {epoch}")

        # RESET/INIT THE METRIC LOGGERS
        self._reset_metrics()

        self.train_metrics.to(self.device)
        loss_avg_meter = core_utils.utils.AverageMeter()

        context = PhaseContext(epoch=epoch,
                               optimizer=self.optimizer,
                               metrics_compute_fn=self.train_metrics,
                               loss_avg_meter=loss_avg_meter,
                               criterion=self.criterion,
                               device=self.device,
                               lr_warmup_epochs=self.training_params.lr_warmup_epochs,
                               sg_logger=self.sg_logger,
                               train_loader=self.train_loader,
                               context_methods=self._get_context_methods(Phase.TRAIN_BATCH_END),
                               ddp_silent_mode=self.ddp_silent_mode)

        for batch_idx, batch_items in enumerate(progress_bar_train_loader):
            batch_items = core_utils.tensor_container_to_device(batch_items, self.device, non_blocking=True)
            inputs, targets, additional_batch_items = sg_trainer_utils.unpack_batch_items(batch_items)

            if self.pre_prediction_callback is not None:
                inputs, targets = self.pre_prediction_callback(inputs, targets, batch_idx)
            # AUTOCAST IS ENABLED ONLY IF self.training_params.mixed_precision - IF enabled=False AUTOCAST HAS NO EFFECT
            with autocast(enabled=self.training_params.mixed_precision):
                # FORWARD PASS TO GET NETWORK'S PREDICTIONS
                outputs = self.net(inputs)

                # COMPUTE THE LOSS FOR BACK PROP + EXTRA METRICS COMPUTED DURING THE LOSS FORWARD PASS
                loss, loss_log_items = self._get_losses(outputs, targets)

            context.update_context(batch_idx=batch_idx,
                                   inputs=inputs,
                                   preds=outputs,
                                   target=targets,
                                   loss_log_items=loss_log_items,
                                   **additional_batch_items)

            self.phase_callback_handler(Phase.TRAIN_BATCH_END, context)

            # LOG LR THAT WILL BE USED IN CURRENT EPOCH AND AFTER FIRST WARMUP/LR_SCHEDULER UPDATE BEFORE WEIGHT UPDATE
            if not self.ddp_silent_mode and batch_idx == 0:
                self._write_lrs(epoch)

            self._backward_step(loss, epoch, batch_idx, context)

            # COMPUTE THE RUNNING USER METRICS AND LOSS RUNNING ITEMS. RESULT TUPLE IS THEIR CONCATENATION.
            logging_values = loss_avg_meter.average + get_metrics_results_tuple(self.train_metrics)
            gpu_memory_utilization = torch.cuda.memory_cached() / 1E9 if torch.cuda.is_available() else 0

            # RENDER METRICS PROGRESS
            pbar_message_dict = get_train_loop_description_dict(logging_values,
                                                                self.train_metrics,
                                                                self.loss_logging_items_names,
                                                                gpu_mem=gpu_memory_utilization)

            progress_bar_train_loader.set_postfix(**pbar_message_dict)

            # TODO: ITERATE BY MAX ITERS
            # FOR INFINITE SAMPLERS WE MUST BREAK WHEN REACHING LEN ITERATIONS.
            if self._infinite_train_loader and batch_idx == len(self.train_loader) - 1:
                break

        if not self.ddp_silent_mode:
            self.sg_logger.upload()

        self.train_monitored_values = sg_trainer_utils.update_monitored_values_dict(
            monitored_values_dict=self.train_monitored_values, new_values_dict=pbar_message_dict)

        return logging_values

    def _get_losses(self, outputs: torch.Tensor, targets: torch.Tensor) -> Tuple[torch.Tensor, tuple]:
        # GET THE OUTPUT OF THE LOSS FUNCTION
        loss = self.criterion(outputs, targets)
        if isinstance(loss, tuple):
            loss, loss_logging_items = loss
            # IF ITS NOT A TUPLE THE LOGGING ITEMS CONTAIN ONLY THE LOSS FOR BACKPROP (USER DEFINED LOSS RETURNS SCALAR)
        else:
            loss_logging_items = loss.unsqueeze(0).detach()

        if len(loss_logging_items) != len(self.loss_logging_items_names):
            raise ValueError("Loss output length must match loss_logging_items_names. Got " + str(
                len(loss_logging_items)) + ', and ' + str(len(self.loss_logging_items_names)))
        # RETURN AND THE LOSS LOGGING ITEMS COMPUTED DURING LOSS FORWARD PASS
        return loss, loss_logging_items

    def _backward_step(self, loss: torch.Tensor, epoch: int, batch_idx: int, context: PhaseContext, *args, **kwargs):
        """
        Run backprop on the loss and perform a step
        :param loss: The value computed by the loss function
        :param optimizer: An object that can perform a gradient step and zeroize model gradient
        :param epoch: number of epoch the training is on
        :param batch_idx: number of iteration inside the current epoch
        :param context: current phase context
        :return:
        """
        # SCALER IS ENABLED ONLY IF self.training_params.mixed_precision=True
        self.scaler.scale(loss).backward()

        # APPLY GRADIENT CLIPPING IF REQUIRED
        if self.training_params.clip_grad_norm:
            torch.nn.utils.clip_grad_norm_(self.net.parameters(), self.training_params.clip_grad_norm)

        # ACCUMULATE GRADIENT FOR X BATCHES BEFORE OPTIMIZING
        integrated_batches_num = batch_idx + len(self.train_loader) * epoch + 1

        if integrated_batches_num % self.batch_accumulate == 0:
            # SCALER IS ENABLED ONLY IF self.training_params.mixed_precision=True
            self.scaler.step(self.optimizer)
            self.scaler.update()

            self.optimizer.zero_grad()
            if self.ema:
                self.ema_model.update(self.net, integrated_batches_num / (len(self.train_loader) * self.max_epochs))

            # RUN PHASE CALLBACKS
            self.phase_callback_handler(Phase.TRAIN_BATCH_STEP, context)

    def _save_checkpoint(self, optimizer=None, epoch: int = None, validation_results_tuple: tuple = None,
                         context: PhaseContext = None):
        """
        Save the current state dict as latest (always), best (if metric was improved), epoch# (if determined in training
        params)
        """
        # WHEN THE validation_results_tuple IS NONE WE SIMPLY SAVE THE state_dict AS LATEST AND Return
        if validation_results_tuple is None:
            self.sg_logger.add_checkpoint(tag='ckpt_latest_weights_only.pth', state_dict={'net': self.net.state_dict()},
                                          global_step=epoch)
            return

        # COMPUTE THE CURRENT metric
        # IF idx IS A LIST - SUM ALL THE VALUES STORED IN THE LIST'S INDICES
        metric = validation_results_tuple[self.metric_idx_in_results_tuple] if isinstance(
            self.metric_idx_in_results_tuple, int) else \
            sum([validation_results_tuple[idx] for idx in self.metric_idx_in_results_tuple])

        # BUILD THE state_dict
        state = {'net': self.net.state_dict(), 'acc': metric, 'epoch': epoch}
        if optimizer is not None:
            state['optimizer_state_dict'] = optimizer.state_dict()

        if self.scaler is not None:
            state['scaler_state_dict'] = self.scaler.state_dict()

        if self.ema:
            state['ema_net'] = self.ema_model.ema.state_dict()
        # SAVES CURRENT MODEL AS ckpt_latest
        self.sg_logger.add_checkpoint(tag='ckpt_latest.pth', state_dict=state, global_step=epoch)

        # SAVE MODEL AT SPECIFIC EPOCHS DETERMINED BY save_ckpt_epoch_list
        if epoch in self.training_params.save_ckpt_epoch_list:
            self.sg_logger.add_checkpoint(tag=f'ckpt_epoch_{epoch}.pth', state_dict=state, global_step=epoch)

        # OVERRIDE THE BEST CHECKPOINT AND best_metric IF metric GOT BETTER THAN THE PREVIOUS BEST
        if (metric > self.best_metric and self.greater_metric_to_watch_is_better) or (
                metric < self.best_metric and not self.greater_metric_to_watch_is_better):
            # STORE THE CURRENT metric AS BEST
            self.best_metric = metric
            self._save_best_checkpoint(epoch, state)

            # RUN PHASE CALLBACKS
            self.phase_callback_handler(Phase.VALIDATION_END_BEST_EPOCH, context)

            if isinstance(metric, torch.Tensor):
                metric = metric.item()
            logger.info("Best checkpoint overriden: validation " + self.metric_to_watch + ": " + str(metric))

        if self.training_params.average_best_models:
            net_for_averaging = self.ema_model.ema if self.ema else self.net
            averaged_model_sd = self.model_weight_averaging.get_average_model(net_for_averaging,
                                                                              validation_results_tuple=validation_results_tuple)
            self.sg_logger.add_checkpoint(tag=self.average_model_checkpoint_filename,
                                          state_dict={'net': averaged_model_sd}, global_step=epoch)

    def _save_best_checkpoint(self, epoch, state):
        self.sg_logger.add_checkpoint(tag=self.ckpt_best_name, state_dict=state, global_step=epoch)

    # FIXME - we need to resolve flake8's 'function is too complex' for this function
    def train(self, training_params: dict = dict()):  # noqa: C901
        """

        train - Trains the Model

        IMPORTANT NOTE: Additional batch parameters can be added as a third item (optional) if a tuple is returned by
          the data loaders, as dictionary. The phase context will hold the additional items, under an attribute with
          the same name as the key in this dictionary. Then such items can be accessed through phase callbacks.


            :param training_params:
                - `max_epochs` : int

                    Number of epochs to run training.

                - `lr_updates` : list(int)

                    List of fixed epoch numbers to perform learning rate updates when `lr_mode='step'`.

                - `lr_decay_factor` : float

                    Decay factor to apply to the learning rate at each update when `lr_mode='step'`.


                -  `lr_mode` : str

                    Learning rate scheduling policy, one of ['step','poly','cosine','function']. 'step' refers to
                    constant updates at epoch numbers passed through `lr_updates`. 'cosine' refers to Cosine Anealing
                    policy as mentioned in https://arxiv.org/abs/1608.03983. 'poly' refers to polynomial decrease i.e
                    in each epoch iteration `self.lr = self.initial_lr * pow((1.0 - (current_iter / max_iter)),
                    0.9)` 'function' refers to user defined learning rate scheduling function, that is passed through
                    `lr_schedule_function`.

                - `lr_schedule_function` : Union[callable,None]

                    Learning rate scheduling function to be used when `lr_mode` is 'function'.

                - `lr_warmup_epochs` : int (default=0)

                    Number of epochs for learning rate warm up - see https://arxiv.org/pdf/1706.02677.pdf (Section 2.2).

                - `cosine_final_lr_ratio` : float (default=0.01)
                    Final learning rate ratio (only relevant when `lr_mode`='cosine'). The cosine starts from initial_lr and reaches
                     initial_lr * cosine_final_lr_ratio in last epoch

                - `inital_lr` : float

                    Initial learning rate.

                - `loss` : Union[nn.module, str]

                    Loss function for training.
                    One of SuperGradient's built in options:

                              "cross_entropy": LabelSmoothingCrossEntropyLoss,
                              "mse": MSELoss,
                              "r_squared_loss": RSquaredLoss,
                              "detection_loss": YoLoV3DetectionLoss,
                              "shelfnet_ohem_loss": ShelfNetOHEMLoss,
                              "shelfnet_se_loss": ShelfNetSemanticEncodingLoss,
                              "ssd_loss": SSDLoss,


                    or user defined nn.module loss function.

                    IMPORTANT: forward(...) should return a (loss, loss_items) tuple where loss is the tensor used
                    for backprop (i.e what your original loss function returns), and loss_items should be a tensor of
                    shape (n_items), of values computed during the forward pass which we desire to log over the
                    entire epoch. For example- the loss itself should always be logged. Another example is a scenario
                    where the computed loss is the sum of a few components we would like to log- these entries in
                    loss_items).

                    When training, set the loss_logging_items_names parameter in train_params to be a list of
                    strings, of length n_items who's ith element is the name of the ith entry in loss_items. Then
                    each item will be logged, rendered on tensorboard and "watched" (i.e saving model checkpoints
                    according to it).

                    Since running logs will save the loss_items in some internal state, it is recommended that
                    loss_items are detached from their computational graph for memory efficiency.

                - `optimizer` : Union[str, torch.optim.Optimizer]

                    Optimization algorithm. One of ['Adam','SGD','RMSProp'] corresponding to the torch.optim
                    optimzers implementations, or any object that implements torch.optim.Optimizer.

                - `criterion_params` : dict

                    Loss function parameters.

                - `optimizer_params` : dict
                    When `optimizer` is one of ['Adam','SGD','RMSProp'], it will be initialized with optimizer_params.

                    (see https://pytorch.org/docs/stable/optim.html for the full list of
                    parameters for each optimizer).

                - `train_metrics_list` : list(torchmetrics.Metric)

                    Metrics to log during training. For more information on torchmetrics see
                    https://torchmetrics.rtfd.io/en/latest/.


                - `valid_metrics_list` : list(torchmetrics.Metric)

                    Metrics to log during validation/testing. For more information on torchmetrics see
                    https://torchmetrics.rtfd.io/en/latest/.


                - `loss_logging_items_names` : list(str)

                    The list of names/titles for the outputs returned from the loss functions forward pass (reminder-
                    the loss function should return the tuple (loss, loss_items)). These names will be used for
                    logging their values.

                - `metric_to_watch` : str (default="Accuracy")

                    will be the metric which the model checkpoint will be saved according to, and can be set to any
                    of the following:

                        a metric name (str) of one of the metric objects from the valid_metrics_list

                        a "metric_name" if some metric in valid_metrics_list has an attribute component_names which
                        is a list referring to the names of each entry in the output metric (torch tensor of size n)

                        one of "loss_logging_items_names" i.e which will correspond to an item returned during the
                        loss function's forward pass.

                    At the end of each epoch, if a new best metric_to_watch value is achieved, the models checkpoint
                    is saved in YOUR_PYTHON_PATH/checkpoints/ckpt_best.pth

                - `greater_metric_to_watch_is_better` : bool

                    When choosing a model's checkpoint to be saved, the best achieved model is the one that maximizes the
                     metric_to_watch when this parameter is set to True, and a one that minimizes it otherwise.

                - `ema` : bool (default=False)

                    Whether to use Model Exponential Moving Average (see
                    https://github.com/rwightman/pytorch-image-models ema implementation)

                - `batch_accumulate` : int (default=1)

                    Number of batches to accumulate before every backward pass.

                - `ema_params` : dict

                    Parameters for the ema model.

                - `zero_weight_decay_on_bias_and_bn` : bool (default=False)

                    Whether to apply weight decay on batch normalization parameters or not (ignored when the passed
                    optimizer has already been initialized).


                - `load_opt_params` : bool (default=True)

                    Whether to load the optimizers parameters as well when loading a model's checkpoint.

                - `run_validation_freq` : int (default=1)

                    The frequency in which validation is performed during training (i.e the validation is ran every
                     `run_validation_freq` epochs.

                - `save_model` : bool (default=True)

                    Whether to save the model checkpoints.

                - `silent_mode` : bool

                    Silents the print outs.

                - `mixed_precision` : bool

                    Whether to use mixed precision or not.

                - `save_ckpt_epoch_list` : list(int) (default=[])

                    List of fixed epoch indices the user wishes to save checkpoints in.

                - `average_best_models` : bool (default=False)

                    If set, a snapshot dictionary file and the average model will be saved / updated at every epoch
                    and evaluated only when training is completed. The snapshot file will only be deleted upon
                    completing the training. The snapshot dict will be managed on cpu.

                - `precise_bn` : bool (default=False)

                    Whether to use precise_bn calculation during the training.

                - `precise_bn_batch_size` : int (default=None)

                    The effective batch size we want to calculate the batchnorm on. For example, if we are training a model
                    on 8 gpus, with a batch of 128 on each gpu, a good rule of thumb would be to give it 8192
                    (ie: effective_batch_size * num_gpus = batch_per_gpu * num_gpus * num_gpus).
                    If precise_bn_batch_size is not provided in the training_params, the latter heuristic will be taken.

                - `seed` : int (default=42)

                    Random seed to be set for torch, numpy, and random. When using DDP each process will have it's seed
                    set to seed + rank.


                - `log_installed_packages` : bool (default=False)

                    When set, the list of all installed packages (and their versions) will be written to the tensorboard
                     and logfile (useful when trying to reproduce results).

                - `dataset_statistics` : bool (default=False)

                    Enable a statistic analysis of the dataset. If set to True the dataset will be analyzed and a report
                    will be added to the tensorboard along with some sample images from the dataset. Currently only
                    detection datasets are supported for analysis.

                -  `save_full_train_log` : bool (default=False)

                    When set, a full log (of all super_gradients modules, including uncaught exceptions from any other
                     module) of the training will be saved in the checkpoint directory under full_train_log.log

                -  `sg_logger` : Union[AbstractSGLogger, str] (defauls=base_sg_logger)

                    Define the SGLogger object for this training process. The SGLogger handles all disk writes, logs, TensorBoard, remote logging
                    and remote storage. By overriding the default base_sg_logger, you can change the storage location, support external monitoring and logging
                    or support remote storage.

                -   `sg_logger_params` : dict

                    SGLogger parameters

                -   `clip_grad_norm` : float

                    Defines a maximal L2 norm of the gradients. Values which exceed the given value will be clipped

                -   `lr_cooldown_epochs` : int (default=0)

                    Number of epochs to cooldown LR (i.e the last epoch from scheduling view point=max_epochs-cooldown).

                -   `pre_prediction_callback` : Callable (default=None)

                     When not None, this callback will be applied to images and targets, and returning them to be used
                      for the forward pass, and further computations. Args for this callable should be in the order
                      (inputs, targets, batch_idx) returning modified_inputs, modified_targets

                -   `ckpt_best_name` : str (default='ckpt_best.pth')

                    The best checkpoint (according to metric_to_watch) will be saved under this filename in the checkpoints directory.

                -   `enable_qat`: bool (default=False)

                    Adds a QATCallback to the phase callbacks, that triggers quantization aware training starting from
                     qat_params["start_epoch"]

                -   `qat_params`: dict-like object with the following key/values:

                        start_epoch: int, first epoch to start QAT.

                        quant_modules_calib_method: str, One of [percentile, mse, entropy, max]. Statistics method for amax
                         computation of the quantized modules (default=percentile).

                        per_channel_quant_modules: bool, whether quant modules should be per channel (default=False).

                        calibrate: bool, whether to perfrom calibration (default=False).

                        calibrated_model_path: str, path to a calibrated checkpoint (default=None).

                        calib_data_loader: torch.utils.data.DataLoader, data loader of the calibration dataset. When None,
                         context.train_loader will be used (default=None).

                        num_calib_batches: int, number of batches to collect the statistics from.

                        percentile: float, percentile value to use when Trainer,quant_modules_calib_method='percentile'.
                         Discarded when other methods are used (Default=99.99).


        :return:
        """
        global logger

        if self.net is None:
            raise Exception('Model', 'No model found')
        if self.dataset_interface is None and self.train_loader is None:
            raise Exception('Data', 'No dataset found')

        self.training_params = TrainingParams()
        self.training_params.override(**training_params)

        # SET RANDOM SEED
        random_seed(is_ddp=self.multi_gpu == MultiGPUMode.DISTRIBUTED_DATA_PARALLEL,
                    device=self.device, seed=self.training_params.seed)

        silent_mode = self.training_params.silent_mode or self.ddp_silent_mode
        # METRICS
        self._set_train_metrics(train_metrics_list=self.training_params.train_metrics_list)
        self._set_valid_metrics(valid_metrics_list=self.training_params.valid_metrics_list)
        self.loss_logging_items_names = self.training_params.loss_logging_items_names

        self.results_titles = ["Train_" + t for t in
                               self.loss_logging_items_names + get_metrics_titles(self.train_metrics)] + \
                              ["Valid_" + t for t in
                               self.loss_logging_items_names + get_metrics_titles(self.valid_metrics)]

        # Store the metric to follow (loss\accuracy) and initialize as the worst value
        self.metric_to_watch = self.training_params.metric_to_watch
        self.greater_metric_to_watch_is_better = self.training_params.greater_metric_to_watch_is_better
        self.metric_idx_in_results_tuple = (self.loss_logging_items_names + get_metrics_titles(self.valid_metrics)).index(self.metric_to_watch)

        # Instantiate the values to monitor (loss/metric)
        for loss in self.loss_logging_items_names:
            self.train_monitored_values[loss] = MonitoredValue(name=loss, greater_is_better=False)
            self.valid_monitored_values[loss] = MonitoredValue(name=loss, greater_is_better=False)
        self.valid_monitored_values[self.metric_to_watch] = MonitoredValue(name=self.metric_to_watch,
                                                                           greater_is_better=True)

        # Allowing loading instantiated loss or string
        if isinstance(self.training_params.loss, str):
            criterion_cls = LOSSES[self.training_params.loss]
            self.criterion = criterion_cls(**self.training_params.criterion_params)

        elif isinstance(self.training_params.loss, Mapping):
            self.criterion = LossesFactory().get(self.training_params.loss)

        elif isinstance(self.training_params.loss, nn.Module):
            self.criterion = self.training_params.loss

        self.criterion.to(self.device)

        self.max_epochs = self.training_params.max_epochs

        self.ema = self.training_params.ema

        self.precise_bn = self.training_params.precise_bn
        self.precise_bn_batch_size = self.training_params.precise_bn_batch_size

        self.batch_accumulate = self.training_params.batch_accumulate
        num_batches = len(self.train_loader)

        if self.ema:
            ema_params = self.training_params.ema_params
            logger.info(f'Using EMA with params {ema_params}')
            self.ema_model = self._instantiate_ema_model(**ema_params)
            self.ema_model.updates = self.start_epoch * num_batches // self.batch_accumulate
            if self.load_checkpoint:
                if 'ema_net' in self.checkpoint.keys():
                    self.ema_model.ema.load_state_dict(self.checkpoint['ema_net'])
                else:
                    self.ema = False
                    logger.warning(
                        "[Warning] Checkpoint does not include EMA weights, continuing training without EMA.")

        self.run_validation_freq = self.training_params.run_validation_freq
        validation_results_tuple = (0, 0)
        inf_time = 0
        timer = core_utils.Timer(self.device)

        # IF THE LR MODE IS NOT DEFAULT TAKE IT FROM THE TRAINING PARAMS
        self.lr_mode = self.training_params.lr_mode
        load_opt_params = self.training_params.load_opt_params

        self.phase_callbacks = self.training_params.phase_callbacks or []
        self.phase_callbacks = ListFactory(CallbacksFactory()).get(self.phase_callbacks)

        if self.lr_mode is not None:
            sg_lr_callback_cls = LR_SCHEDULERS_CLS_DICT[self.lr_mode]
            self.phase_callbacks.append(sg_lr_callback_cls(train_loader_len=len(self.train_loader),
                                                           net=self.net,
                                                           training_params=self.training_params,
                                                           update_param_groups=self.update_param_groups,
                                                           **self.training_params.to_dict()))
        if self.training_params.lr_warmup_epochs > 0:
            warmup_mode = self.training_params.warmup_mode
            if isinstance(warmup_mode, str):
                warmup_callback_cls = LR_WARMUP_CLS_DICT[warmup_mode]
            elif isinstance(warmup_mode, type) and issubclass(warmup_mode, LRCallbackBase):
                warmup_callback_cls = warmup_mode
            else:
                raise RuntimeError('warmup_mode has to be either a name of a mode (str) or a subclass of PhaseCallback')
            self.phase_callbacks.append(warmup_callback_cls(train_loader_len=len(self.train_loader),
                                                            net=self.net,
                                                            training_params=self.training_params,
                                                            update_param_groups=self.update_param_groups,
                                                            **self.training_params.to_dict()))

        self._add_metrics_update_callback(Phase.TRAIN_BATCH_END)
        self._add_metrics_update_callback(Phase.VALIDATION_BATCH_END)

        # ADD CALLBACK FOR QAT
        self.enable_qat = core_utils.get_param(self.training_params, "enable_qat", False)
        if self.enable_qat:
            self.qat_params = core_utils.get_param(self.training_params, "qat_params")
            if self.qat_params is None:
                raise ValueError("Must pass QAT params when enable_qat=True")
            self.phase_callbacks.append(QATCallback(**self.qat_params))

        self.phase_callback_handler = CallbackHandler(callbacks=self.phase_callbacks)

        if not self.ddp_silent_mode:
            self._initialize_sg_logger_objects()

            if self.training_params.dataset_statistics:
                dataset_statistics_logger = DatasetStatisticsTensorboardLogger(self.sg_logger)
                dataset_statistics_logger.analyze(self.train_loader, all_classes=self.classes,
                                                  title="Train-set", anchors=self.net.module.arch_params.anchors)
                dataset_statistics_logger.analyze(self.valid_loader, all_classes=self.classes,
                                                  title="val-set")
            # AVERAGE BEST 10 MODELS PARAMS
            if self.training_params.average_best_models:
                self.model_weight_averaging = ModelWeightAveraging(self.checkpoints_dir_path,
                                                                   greater_is_better=self.greater_metric_to_watch_is_better,
                                                                   source_ckpt_folder_name=self.source_ckpt_folder_name,
                                                                   metric_to_watch=self.metric_to_watch,
                                                                   metric_idx=self.metric_idx_in_results_tuple,
                                                                   load_checkpoint=self.load_checkpoint,
                                                                   model_checkpoints_location=self.model_checkpoints_location)
        if self.training_params.save_full_train_log and not self.ddp_silent_mode:
            logger = get_logger(__name__,
                                training_log_path=self.sg_logger.log_file_path.replace('.txt', 'full_train_log.log'))
            sg_trainer_utils.log_uncaught_exceptions(logger)

        if not self.load_checkpoint or self.load_weights_only:
            # WHEN STARTING TRAINING FROM SCRATCH, DO NOT LOAD OPTIMIZER PARAMS (EVEN IF LOADING BACKBONE)
            self.start_epoch = 0
            self._reset_best_metric()
            load_opt_params = False

        if isinstance(self.training_params.optimizer, str) or \
                (inspect.isclass(self.training_params.optimizer) and issubclass(self.training_params.optimizer, torch.optim.Optimizer)):
            self.optimizer = build_optimizer(net=self.net, lr=self.training_params.initial_lr,
                                             training_params=self.training_params)
        elif isinstance(self.training_params.optimizer, torch.optim.Optimizer):
            self.optimizer = self.training_params.optimizer
        else:
            raise UnsupportedOptimizerFormat()

        # VERIFY GRADIENT CLIPPING VALUE
        if self.training_params.clip_grad_norm is not None and self.training_params.clip_grad_norm <= 0:
            raise TypeError('Params', 'Invalid clip_grad_norm')

        if self.load_checkpoint and load_opt_params:
            self.optimizer.load_state_dict(self.checkpoint['optimizer_state_dict'])

        self.pre_prediction_callback = CallbacksFactory().get(self.training_params.pre_prediction_callback)

        self._initialize_mixed_precision(self.training_params.mixed_precision)

        self._infinite_train_loader = (hasattr(self.train_loader, "sampler") and isinstance(self.train_loader.sampler, InfiniteSampler)) or \
                                      (hasattr(self.train_loader, "batch_sampler") and isinstance(self.train_loader.batch_sampler.sampler, InfiniteSampler))

        self.ckpt_best_name = self.training_params.ckpt_best_name

        context = PhaseContext(optimizer=self.optimizer,
                               net=self.net,
                               experiment_name=self.experiment_name,
                               ckpt_dir=self.checkpoints_dir_path,
                               criterion=self.criterion,
                               lr_warmup_epochs=self.training_params.lr_warmup_epochs,
                               sg_logger=self.sg_logger,
                               train_loader=self.train_loader,
                               valid_loader=self.valid_loader,
                               training_params=self.training_params,
                               ddp_silent_mode=self.ddp_silent_mode,
                               checkpoint_params=self.checkpoint_params,
                               architecture=self.architecture,
                               arch_params=self.arch_params,
                               metric_idx_in_results_tuple=self.metric_idx_in_results_tuple,
                               metric_to_watch=self.metric_to_watch,
                               device=self.device,
                               context_methods=self._get_context_methods(Phase.PRE_TRAINING)
                               )

        self.phase_callback_handler(Phase.PRE_TRAINING, context)

        try:
            # HEADERS OF THE TRAINING PROGRESS
            if not silent_mode:
                logger.info(
                    f'Started training for {self.max_epochs - self.start_epoch} epochs ({self.start_epoch}/'f'{self.max_epochs - 1})\n')
            for epoch in range(self.start_epoch, self.max_epochs):
                if context.stop_training:
                    logger.info("Request to stop training has been received, stopping training")
                    break

                # Phase.TRAIN_EPOCH_START
                # RUN PHASE CALLBACKS
                context.update_context(epoch=epoch)
                self.phase_callback_handler(Phase.TRAIN_EPOCH_START, context)

                # IN DDP- SET_EPOCH WILL CAUSE EVERY PROCESS TO BE EXPOSED TO THE ENTIRE DATASET BY SHUFFLING WITH A
                # DIFFERENT SEED EACH EPOCH START
                if self.multi_gpu == MultiGPUMode.DISTRIBUTED_DATA_PARALLEL and hasattr(self.train_loader, "sampler") \
                        and hasattr(self.train_loader.sampler, "set_epoch"):
                    self.train_loader.sampler.set_epoch(epoch)

                train_metrics_tuple = self._train_epoch(epoch=epoch, silent_mode=silent_mode)

                # Phase.TRAIN_EPOCH_END
                # RUN PHASE CALLBACKS
                train_metrics_dict = get_metrics_dict(train_metrics_tuple, self.train_metrics,
                                                      self.loss_logging_items_names)

                context.update_context(metrics_dict=train_metrics_dict)
                self.phase_callback_handler(Phase.TRAIN_EPOCH_END, context)

                # CALCULATE PRECISE BATCHNORM STATS
                if self.precise_bn:
                    compute_precise_bn_stats(model=self.net, loader=self.train_loader,
                                             precise_bn_batch_size=self.precise_bn_batch_size,
                                             num_gpus=self.num_devices)
                    if self.ema:
                        compute_precise_bn_stats(model=self.ema_model.ema, loader=self.train_loader,
                                                 precise_bn_batch_size=self.precise_bn_batch_size,
                                                 num_gpus=self.num_devices)

                # model switch - we replace self.net.module with the ema model for the testing and saving part
                # and then switch it back before the next training epoch
                if self.ema:
                    self.ema_model.update_attr(self.net)
                    keep_model = self.net
                    self.net = self.ema_model.ema

                # RUN TEST ON VALIDATION SET EVERY self.run_validation_freq EPOCHS
                if (epoch + 1) % self.run_validation_freq == 0:
                    timer.start()
                    validation_results_tuple = self._validate_epoch(epoch=epoch, silent_mode=silent_mode)
                    inf_time = timer.stop()

                    # Phase.VALIDATION_EPOCH_END
                    # RUN PHASE CALLBACKS
                    valid_metrics_dict = get_metrics_dict(validation_results_tuple, self.valid_metrics,
                                                          self.loss_logging_items_names)

                    context.update_context(metrics_dict=valid_metrics_dict)
                    self.phase_callback_handler(Phase.VALIDATION_EPOCH_END, context)

                if self.ema:
                    self.net = keep_model

                if not self.ddp_silent_mode:
                    # SAVING AND LOGGING OCCURS ONLY IN THE MAIN PROCESS (IN CASES THERE ARE SEVERAL PROCESSES - DDP)
                    self._write_to_disk_operations(train_metrics_tuple, validation_results_tuple, inf_time, epoch,
                                                   context)

            # Evaluating the average model and removing snapshot averaging file if training is completed
            if self.training_params.average_best_models:
                self._validate_final_average_model(cleanup_snapshots_pkl_file=True)

        except KeyboardInterrupt:
            logger.info(
                '\n[MODEL TRAINING EXECUTION HAS BEEN INTERRUPTED]... Please wait until SOFT-TERMINATION process '
                'finishes and saves all of the Model Checkpoints and log files before terminating...')
            logger.info('For HARD Termination - Stop the process again')

        finally:
            if self.multi_gpu == MultiGPUMode.DISTRIBUTED_DATA_PARALLEL:
                # CLEAN UP THE MULTI-GPU PROCESS GROUP WHEN DONE
                if torch.distributed.is_initialized():
                    torch.distributed.destroy_process_group()

            # PHASE.TRAIN_END
            self.phase_callback_handler(Phase.POST_TRAINING, context)

            if not self.ddp_silent_mode:
                if self.model_checkpoints_location != 'local':
                    logger.info('[CLEANUP] - Saving Checkpoint files')
                    self.sg_logger.upload()

                self.sg_logger.close()

    def _reset_best_metric(self):
        self.best_metric = -1 * np.inf if self.greater_metric_to_watch_is_better else np.inf

    def _reset_metrics(self):
        for metric in ("train_metrics", "valid_metrics", "test_metrics"):
            if hasattr(self, metric) and getattr(self, metric) is not None:
                getattr(self, metric).reset()

    @resolve_param('train_metrics_list', ListFactory(MetricsFactory()))
    def _set_train_metrics(self, train_metrics_list):
        self.train_metrics = MetricCollection(train_metrics_list)

    @resolve_param('valid_metrics_list', ListFactory(MetricsFactory()))
    def _set_valid_metrics(self, valid_metrics_list):
        self.valid_metrics = MetricCollection(valid_metrics_list)

    def _initialize_mixed_precision(self, mixed_precision_enabled: bool):
        # SCALER IS ALWAYS INITIALIZED BUT IS DISABLED IF MIXED PRECISION WAS NOT SET
        self.scaler = GradScaler(enabled=mixed_precision_enabled)

        if mixed_precision_enabled:
            assert self.device.startswith('cuda'), "mixed precision is not available for CPU"
            if self.multi_gpu == MultiGPUMode.DATA_PARALLEL:
                # IN DATAPARALLEL MODE WE NEED TO WRAP THE FORWARD FUNCTION OF OUR MODEL SO IT WILL RUN WITH AUTOCAST.
                # BUT SINCE THE MODULE IS CLONED TO THE DEVICES ON EACH FORWARD CALL OF A DATAPARALLEL MODEL,
                # WE HAVE TO REGISTER THE WRAPPER BEFORE EVERY FORWARD CALL
                def hook(module, _):
                    module.forward = MultiGPUModeAutocastWrapper(module.forward)

                self.net.module.register_forward_pre_hook(hook=hook)

            if self.load_checkpoint:
                scaler_state_dict = core_utils.get_param(self.checkpoint, 'scaler_state_dict')
                if scaler_state_dict is None:
                    logger.warning(
                        'Mixed Precision - scaler state_dict not found in loaded model. This may case issues '
                        'with loss scaling')
                else:
                    self.scaler.load_state_dict(scaler_state_dict)

    def _validate_final_average_model(self, cleanup_snapshots_pkl_file=False):
        """
        Testing the averaged model by loading the last saved average checkpoint and running test.
        Will be loaded to each of DDP processes
        :param cleanup_pkl_file: a flag for deleting the 10 best snapshots dictionary
        """
        logger.info('RUNNING ADDITIONAL TEST ON THE AVERAGED MODEL...')

        keep_state_dict = deepcopy(self.net.state_dict())
        # SETTING STATE DICT TO THE AVERAGE MODEL FOR EVALUATION
        average_model_ckpt_path = os.path.join(self.checkpoints_dir_path, self.average_model_checkpoint_filename)
        average_model_sd = read_ckpt_state_dict(average_model_ckpt_path)['net']

        self.net.load_state_dict(average_model_sd)
        # testing the averaged model and save instead of best model if needed
        averaged_model_results_tuple = self._validate_epoch(epoch=self.max_epochs)

        # Reverting the current model
        self.net.load_state_dict(keep_state_dict)

        if not self.ddp_silent_mode:
            # Adding values to sg_logger
            # looping over last titles which corresponds to validation (and average model) metrics.
            all_titles = self.results_titles[-1 * len(averaged_model_results_tuple):]
            result_dict = {all_titles[i]: averaged_model_results_tuple[i] for i in
                           range(len(averaged_model_results_tuple))}

            self.sg_logger.add_scalars(tag_scalar_dict=result_dict, global_step=self.max_epochs)

            average_model_tb_titles = ['Averaged Model ' + x for x in
                                       self.results_titles[-1 * len(averaged_model_results_tuple):]]
            write_struct = ''
            for ind, title in enumerate(average_model_tb_titles):
                write_struct += '%s: %.3f  \n  ' % (title, averaged_model_results_tuple[ind])
                self.sg_logger.add_scalar(title, averaged_model_results_tuple[ind], global_step=self.max_epochs)

            self.sg_logger.add_text("Averaged_Model_Performance", write_struct, self.max_epochs)
            if cleanup_snapshots_pkl_file:
                self.model_weight_averaging.cleanup()

    @property
    def get_arch_params(self):
        return self.arch_params.to_dict()

    @property
    def get_structure(self):
        return self.net.module.structure

    @property
    def get_architecture(self):
        return self.architecture

    def set_experiment_name(self, experiment_name):
        self.experiment_name = experiment_name

    def _re_build_model(self, arch_params={}):
        """
        arch_params : dict
            Architecture H.P. e.g.: block, num_blocks, num_classes, etc.
        :return:
        """
        if 'num_classes' not in arch_params.keys():
            if self.dataset_interface is None:
                raise Exception('Error', 'Number of classes not defined in arch params and dataset is not defined')
            else:
                arch_params['num_classes'] = len(self.classes)

        self.arch_params = core_utils.HpmStruct(**arch_params)
        self.classes = self.arch_params.num_classes
        self.net = self._instantiate_net(self.architecture, self.arch_params, self.checkpoint_params)
        # save the architecture for neural architecture search
        if hasattr(self.net, 'structure'):
            self.architecture = self.net.structure

        self.net.to(self.device)

        if self.multi_gpu == MultiGPUMode.DISTRIBUTED_DATA_PARALLEL:
            logger.warning("Warning: distributed training is not supported in re_build_model()")
        self.net = torch.nn.DataParallel(self.net,
                                         device_ids=self.device_ids) if self.multi_gpu else core_utils.WrappedModel(
            self.net)

    @property
    def get_module(self):
        return self.net

    def set_module(self, module):
        self.net = module

    def _initialize_device(self, requested_device: str, requested_multi_gpu: Union[MultiGPUMode, str]):
        """
        _initialize_device - Initializes the device for the model - Default is CUDA
            :param requested_device:        Device to initialize ('cuda' / 'cpu')
            :param requested_multi_gpu:     Get Multiple GPU
        """

        if isinstance(requested_multi_gpu, str):
            requested_multi_gpu = MultiGPUMode(requested_multi_gpu)

        # SELECT CUDA DEVICE
        if requested_device == 'cuda':
            if torch.cuda.is_available():
                self.device = 'cuda'  # TODO - we may want to set the device number as well i.e. 'cuda:1'
            else:
                raise RuntimeError('CUDA DEVICE NOT FOUND... EXITING')

        # SELECT CPU DEVICE
        elif requested_device == 'cpu':
            self.device = 'cpu'
            self.multi_gpu = False
        else:
            # SELECT CUDA DEVICE BY DEFAULT IF AVAILABLE
            self.device = 'cuda' if torch.cuda.is_available() else 'cpu'

        # DEFUALT IS SET TO 1 - IT IS CHANGED IF MULTI-GPU IS USED
        self.num_devices = 1

        # IN CASE OF MULTIPLE GPUS UPDATE THE LEARNING AND DATA PARAMETERS
        # FIXME - CREATE A DISCUSSION ON THESE PARAMETERS - WE MIGHT WANT TO CHANGE THE WAY WE USE THE LR AND
        if requested_multi_gpu != MultiGPUMode.OFF:
            if 'cuda' in self.device:
                # COLLECT THE AVAILABLE GPU AND COUNT THE AVAILABLE GPUS AMOUNT
                self.device_ids = list(range(torch.cuda.device_count()))
                self.num_devices = len(self.device_ids)
                if self.num_devices == 1:
                    self.multi_gpu = MultiGPUMode.OFF
                    if requested_multi_gpu != MultiGPUMode.AUTO:
                        # if AUTO mode was set - do not log a warning
                        logger.warning(
                            '\n[WARNING] - Tried running on multiple GPU but only a single GPU is available\n')
                else:
                    if requested_multi_gpu == MultiGPUMode.AUTO:
                        if env_helpers.is_distributed():
                            requested_multi_gpu = MultiGPUMode.DISTRIBUTED_DATA_PARALLEL
                        else:
                            requested_multi_gpu = MultiGPUMode.DATA_PARALLEL

                    self.multi_gpu = requested_multi_gpu
                    if self.multi_gpu == MultiGPUMode.DISTRIBUTED_DATA_PARALLEL:
                        self._initialize_ddp()
            else:
                # MULTIPLE GPUS CAN BE ACTIVE ONLY IF A GPU IS AVAILABLE
                self.multi_gpu = MultiGPUMode.OFF
                logger.warning('\n[WARNING] - Tried running on multiple GPU but none are available => running on CPU\n')

    def _initialize_ddp(self):
        """
        Initializes Distributed Data Parallel

        Usage:

            python -m torch.distributed.launch --nproc_per_node=n YOUR_TRAINING_SCRIPT.py
            where n is the number of GPUs required, e.g., n=8

            Important note: (1) in distributed training it is customary to specify learning rates and batch sizes per GPU.
            Whatever learning rate and schedule you specify will be applied to the each GPU individually.
            Since gradients are passed and summed (reduced) from all to all GPUs, the effective batch size is the
            batch you specify times the number of GPUs. In the literature there are several "best practices" to set
            learning rates and schedules for large batch sizes.

        """
        logger.info("Distributed training starting...")
        local_rank = environment_config.DDP_LOCAL_RANK
        if not torch.distributed.is_initialized():
            torch.distributed.init_process_group(backend='nccl', init_method='env://')

        if local_rank > 0:
            f = open(os.devnull, 'w')
            sys.stdout = f  # silent all printing for non master process

        torch.cuda.set_device(local_rank)
        self.device = 'cuda:%d' % local_rank

        # MAKE ALL HIGHER-RANK GPUS SILENT (DISTRIBUTED MODE)
        self.ddp_silent_mode = local_rank > 0

        if torch.distributed.get_rank() == 0:
            logger.info(f"Training in distributed mode... with {str(torch.distributed.get_world_size())} GPUs")

    def _switch_device(self, new_device):
        self.device = new_device
        self.net.to(self.device)

    # FIXME - we need to resolve flake8's 'function is too complex' for this function
    def _load_checkpoint_to_model(self):  # noqa: C901 - too complex
        """
        Copies the source checkpoint to a local folder and loads the checkpoint's data to the model using the
         attributes:

         strict:           See StrictLoad class documentation for details.
         load_backbone:    loads the provided checkpoint to self.net.backbone instead of self.net
         source_ckpt_folder_name: The folder where the checkpoint is saved. By default uses the self.experiment_name

        NOTE: 'acc', 'epoch', 'optimizer_state_dict' and the logs are NOT loaded if self.zeroize_prev_train_params
         is True
        """

        self._set_ckpt_loading_attributes()

        if self.load_checkpoint or self.external_checkpoint_path:
            # GET LOCAL PATH TO THE CHECKPOINT FILE FIRST
            ckpt_local_path = get_ckpt_local_path(source_ckpt_folder_name=self.source_ckpt_folder_name,
                                                  experiment_name=self.experiment_name,
                                                  ckpt_name=self.ckpt_name,
                                                  model_checkpoints_location=self.model_checkpoints_location,
                                                  external_checkpoint_path=self.external_checkpoint_path,
                                                  overwrite_local_checkpoint=self.overwrite_local_checkpoint,
                                                  load_weights_only=self.load_weights_only)

            # LOAD CHECKPOINT TO MODEL
            self.checkpoint = load_checkpoint_to_model(ckpt_local_path=ckpt_local_path,
                                                       load_backbone=self.load_backbone,
                                                       net=self.net,
                                                       strict=self.strict_load.value if isinstance(self.strict_load,
                                                                                                   StrictLoad) else self.strict_load,
                                                       load_weights_only=self.load_weights_only,
                                                       load_ema_as_net=self.load_ema_as_net)

            if 'ema_net' in self.checkpoint.keys():
                logger.warning("[WARNING] Main network has been loaded from checkpoint but EMA network exists as "
                               "well. It "
                               " will only be loaded during validation when training with ema=True. ")

        # UPDATE TRAINING PARAMS IF THEY EXIST & WE ARE NOT LOADING AN EXTERNAL MODEL's WEIGHTS
        self.best_metric = self.checkpoint['acc'] if 'acc' in self.checkpoint.keys() else -1
        self.start_epoch = self.checkpoint['epoch'] if 'epoch' in self.checkpoint.keys() else 0

    def _prep_for_test(self, test_loader: torch.utils.data.DataLoader = None, loss=None, post_prediction_callback=None,
                       test_metrics_list=None,
                       loss_logging_items_names=None, test_phase_callbacks=None):
        """Run commands that are common to all SgModels"""
        # SET THE MODEL IN evaluation STATE
        self.net.eval()

        # IF SPECIFIED IN THE FUNCTION CALL - OVERRIDE THE self ARGUMENTS
        self.test_loader = test_loader or self.test_loader
        self.criterion = loss or self.criterion
        self.post_prediction_callback = post_prediction_callback or self.post_prediction_callback
        self.loss_logging_items_names = loss_logging_items_names or self.loss_logging_items_names
        self.phase_callbacks = test_phase_callbacks or self.phase_callbacks

        if self.phase_callbacks is None:
            self.phase_callbacks = []

        if test_metrics_list:
            self.test_metrics = MetricCollection(test_metrics_list)
            self._add_metrics_update_callback(Phase.TEST_BATCH_END)
            self.phase_callback_handler = CallbackHandler(self.phase_callbacks)

        # WHEN TESTING WITHOUT A LOSS FUNCTION- CREATE EPOCH HEADERS FOR PRINTS
        if self.criterion is None:
            self.loss_logging_items_names = []

        if self.test_metrics is None:
            raise ValueError("Metrics are required to perform test. Pass them through test_metrics_list arg when "
                             "calling test or through training_params when calling train(...)")
        if self.test_loader is None:
            raise ValueError("Test dataloader is required to perform test. Make sure to either pass it through "
                             "test_loader arg or calling connect_dataset_interface upon a DatasetInterface instance "
                             "with a non empty testset attribute.")

        # RESET METRIC RUNNERS
        self._reset_metrics()
        self.test_metrics.to(self.device)

    def _add_metrics_update_callback(self, phase: Phase):
        """
        Adds MetricsUpdateCallback to be fired at phase

        :param phase: Phase for the metrics callback to be fired at
        """
        self.phase_callbacks.append(MetricsUpdateCallback(phase))

    def _initialize_sg_logger_objects(self):
        """Initialize object that collect, write to disk, monitor and store remotely all training outputs"""
        sg_logger = core_utils.get_param(self.training_params, 'sg_logger')

        # OVERRIDE SOME PARAMETERS TO MAKE SURE THEY MATCH THE TRAINING PARAMETERS
        general_sg_logger_params = {'experiment_name': self.experiment_name,
                                    'storage_location': self.model_checkpoints_location,
                                    'resumed': self.load_checkpoint,
                                    'training_params': self.training_params,
                                    'checkpoints_dir_path': self.checkpoints_dir_path}

        if sg_logger is None:
            raise RuntimeError('sg_logger must be defined in training params (see default_training_params)')

        if isinstance(sg_logger, AbstractSGLogger):
            self.sg_logger = sg_logger
        elif isinstance(sg_logger, str):
            sg_logger_params = core_utils.get_param(self.training_params, 'sg_logger_params', {})
            if issubclass(SG_LOGGERS[sg_logger], BaseSGLogger):
                sg_logger_params = {**sg_logger_params, **general_sg_logger_params}
            if sg_logger not in SG_LOGGERS:
                raise RuntimeError('sg_logger not defined in SG_LOGGERS')

            self.sg_logger = SG_LOGGERS[sg_logger](**sg_logger_params)
        else:
            raise RuntimeError('sg_logger can be either an sg_logger name (str) or an instance of AbstractSGLogger')

        if not isinstance(self.sg_logger, BaseSGLogger):
            logger.warning("WARNING! Using a user-defined sg_logger: files will not be automatically written to disk!\n"
                           "Please make sure the provided sg_logger writes to disk or compose your sg_logger to BaseSGLogger")

        # IN CASE SG_LOGGER UPDATED THE DIR PATH
        self.checkpoints_dir_path = self.sg_logger.local_dir()
        hyper_param_config = self._get_hyper_param_config()

        self.sg_logger.add_config("hyper_params", hyper_param_config)

        self.sg_logger.flush()

    def _get_hyper_param_config(self):
        """
        Creates a training hyper param config for logging.
        """
        additional_log_items = {'initial_LR': self.training_params.initial_lr,
                                'num_devices': self.num_devices,
                                'multi_gpu': str(self.multi_gpu),
                                'device_type': torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'cpu'}
        # ADD INSTALLED PACKAGE LIST + THEIR VERSIONS
        if self.training_params.log_installed_packages:
            pkg_list = list(map(lambda pkg: str(pkg), _get_installed_distributions()))
            additional_log_items['installed_packages'] = pkg_list
        hyper_param_config = {"arch_params": self.arch_params.__dict__,
                              "checkpoint_params": self.checkpoint_params.__dict__,
                              "training_hyperparams": self.training_params.__dict__,
                              "dataset_params": self.dataset_params.__dict__,
                              "additional_log_items": additional_log_items}
        return hyper_param_config

    def _write_to_disk_operations(self, train_metrics: tuple, validation_results: tuple, inf_time: float, epoch: int,
                                  context: PhaseContext):
        """Run the various logging operations, e.g.: log file, Tensorboard, save checkpoint etc."""
        # STORE VALUES IN A TENSORBOARD FILE
        train_results = list(train_metrics) + list(validation_results) + [inf_time]
        all_titles = self.results_titles + ['Inference Time']

        result_dict = {all_titles[i]: train_results[i] for i in range(len(train_results))}
        self.sg_logger.add_scalars(tag_scalar_dict=result_dict, global_step=epoch)

        # SAVE THE CHECKPOINT
        if self.training_params.save_model:
            self._save_checkpoint(self.optimizer, epoch + 1, validation_results, context)

    def _write_lrs(self, epoch):
        lrs = [self.optimizer.param_groups[i]['lr'] for i in range(len(self.optimizer.param_groups))]
        lr_titles = ['LR/Param_group_' + str(i) for i in range(len(self.optimizer.param_groups))] if len(
            self.optimizer.param_groups) > 1 else ['LR']
        lr_dict = {lr_titles[i]: lrs[i] for i in range(len(lrs))}
        self.sg_logger.add_scalars(tag_scalar_dict=lr_dict, global_step=epoch)

    def test(self,  # noqa: C901
             test_loader: torch.utils.data.DataLoader = None,
             loss: torch.nn.modules.loss._Loss = None,
             silent_mode: bool = False,
             test_metrics_list=None,
             loss_logging_items_names=None, metrics_progress_verbose=False, test_phase_callbacks=None,
             use_ema_net=True) -> tuple:
        """
        Evaluates the model on given dataloader and metrics.

        :param test_loader: dataloader to perform test on.
        :param test_metrics_list: (list(torchmetrics.Metric)) metrics list for evaluation.
        :param silent_mode: (bool) controls verbosity
        :param metrics_progress_verbose: (bool) controls the verbosity of metrics progress (default=False). Slows down the program.
        :param use_ema_net (bool) whether to perform test on self.ema_model.ema (when self.ema_model.ema exists,
            otherwise self.net will be tested) (default=True)
        :return: results tuple (tuple) containing the loss items and metric values.

        All of the above args will override Trainer's corresponding attribute when not equal to None. Then evaluation
         is ran on self.test_loader with self.test_metrics.
        """

        # IN CASE TRAINING WAS PERFROMED BEFORE TEST- MAKE SURE TO TEST THE EMA MODEL (UNLESS SPECIFIED OTHERWISE BY
        # use_ema_net)

        if use_ema_net and self.ema_model is not None:
            keep_model = self.net
            self.net = self.ema_model.ema

        self._prep_for_test(test_loader=test_loader,
                            loss=loss,
                            test_metrics_list=test_metrics_list,
                            loss_logging_items_names=loss_logging_items_names,
                            test_phase_callbacks=test_phase_callbacks,
                            )

        test_results = self.evaluate(data_loader=self.test_loader,
                                     metrics=self.test_metrics,
                                     evaluation_type=EvaluationType.TEST,
                                     silent_mode=silent_mode,
                                     metrics_progress_verbose=metrics_progress_verbose)

        # SWITCH BACK BETWEEN NETS SO AN ADDITIONAL TRAINING CAN BE DONE AFTER TEST
        if use_ema_net and self.ema_model is not None:
            self.net = keep_model

        return test_results

    def _validate_epoch(self, epoch: int, silent_mode: bool = False) -> tuple:
        """
        Runs evaluation on self.valid_loader, with self.valid_metrics.

        :param epoch: (int) epoch idx
        :param silent_mode: (bool) controls verbosity

        :return: results tuple (tuple) containing the loss items and metric values.
        """

        self.net.eval()
        self._reset_metrics()
        self.valid_metrics.to(self.device)

        return self.evaluate(data_loader=self.valid_loader, metrics=self.valid_metrics,
                             evaluation_type=EvaluationType.VALIDATION, epoch=epoch, silent_mode=silent_mode)

    def evaluate(self, data_loader: torch.utils.data.DataLoader, metrics: MetricCollection,
                 evaluation_type: EvaluationType, epoch: int = None, silent_mode: bool = False,
                 metrics_progress_verbose: bool = False):
        """
        Evaluates the model on given dataloader and metrics.

        :param data_loader: dataloader to perform evaluataion on
        :param metrics: (MetricCollection) metrics for evaluation
        :param evaluation_type: (EvaluationType) controls which phase callbacks will be used (for example, on batch end,
            when evaluation_type=EvaluationType.VALIDATION the Phase.VALIDATION_BATCH_END callbacks will be triggered)
        :param epoch: (int) epoch idx
        :param silent_mode: (bool) controls verbosity
        :param metrics_progress_verbose: (bool) controls the verbosity of metrics progress (default=False).
            Slows down the program significantly.

        :return: results tuple (tuple) containing the loss items and metric values.
        """

        # THE DISABLE FLAG CONTROLS WHETHER THE PROGRESS BAR IS SILENT OR PRINTS THE LOGS
        progress_bar_data_loader = tqdm(data_loader, bar_format="{l_bar}{bar:10}{r_bar}", dynamic_ncols=True,
                                        disable=silent_mode)
        loss_avg_meter = core_utils.utils.AverageMeter()
        logging_values = None
        loss_tuple = None
        lr_warmup_epochs = self.training_params.lr_warmup_epochs if self.training_params else None
        context = PhaseContext(epoch=epoch,
                               metrics_compute_fn=metrics,
                               loss_avg_meter=loss_avg_meter,
                               criterion=self.criterion,
                               device=self.device,
                               lr_warmup_epochs=lr_warmup_epochs,
                               sg_logger=self.sg_logger,
                               context_methods=self._get_context_methods(Phase.VALIDATION_BATCH_END))

        if not silent_mode:
            # PRINT TITLES
            pbar_start_msg = f"Validation epoch {epoch}" if evaluation_type == EvaluationType.VALIDATION else "Test"
            progress_bar_data_loader.set_description(pbar_start_msg)

        with torch.no_grad():
            for batch_idx, batch_items in enumerate(progress_bar_data_loader):
                batch_items = core_utils.tensor_container_to_device(batch_items, self.device, non_blocking=True)
                inputs, targets, additional_batch_items = sg_trainer_utils.unpack_batch_items(batch_items)

                output = self.net(inputs)

                if self.criterion is not None:
                    # STORE THE loss_items ONLY, THE 1ST RETURNED VALUE IS THE loss FOR BACKPROP DURING TRAINING
                    loss_tuple = self._get_losses(output, targets)[1].cpu()

                context.update_context(batch_idx=batch_idx,
                                       inputs=inputs,
                                       preds=output,
                                       target=targets,
                                       loss_log_items=loss_tuple,
                                       **additional_batch_items)

                # TRIGGER PHASE CALLBACKS CORRESPONDING TO THE EVALUATION TYPE
                if evaluation_type == EvaluationType.VALIDATION:
                    self.phase_callback_handler(Phase.VALIDATION_BATCH_END, context)
                else:
                    self.phase_callback_handler(Phase.TEST_BATCH_END, context)

                # COMPUTE METRICS IF PROGRESS VERBOSITY IS SET
                if metrics_progress_verbose and not silent_mode:
                    # COMPUTE THE RUNNING USER METRICS AND LOSS RUNNING ITEMS. RESULT TUPLE IS THEIR CONCATENATION.
                    logging_values = get_logging_values(loss_avg_meter, metrics, self.criterion)
                    pbar_message_dict = get_train_loop_description_dict(logging_values,
                                                                        metrics,
                                                                        self.loss_logging_items_names)

                    progress_bar_data_loader.set_postfix(**pbar_message_dict)

        # NEED TO COMPUTE METRICS FOR THE FIRST TIME IF PROGRESS VERBOSITY IS NOT SET
        if not metrics_progress_verbose:
            # COMPUTE THE RUNNING USER METRICS AND LOSS RUNNING ITEMS. RESULT TUPLE IS THEIR CONCATENATION.
            logging_values = get_logging_values(loss_avg_meter, metrics, self.criterion)
            pbar_message_dict = get_train_loop_description_dict(logging_values,
                                                                metrics,
                                                                self.loss_logging_items_names)

            progress_bar_data_loader.set_postfix(**pbar_message_dict)

        # TODO: SUPPORT PRINTING AP PER CLASS- SINCE THE METRICS ARE NOT HARD CODED ANYMORE (as done in
        #  calc_batch_prediction_accuracy_per_class in metric_utils.py), THIS IS ONLY RELEVANT WHEN CHOOSING
        #  DETECTIONMETRICS, WHICH ALREADY RETURN THE METRICS VALUEST HEMSELVES AND NOT THE ITEMS REQUIRED FOR SUCH
        #  COMPUTATION. ALSO REMOVE THE BELOW LINES BY IMPLEMENTING CRITERION AS A TORCHMETRIC.

        if self.multi_gpu == MultiGPUMode.DISTRIBUTED_DATA_PARALLEL:
            logging_values = reduce_results_tuple_for_ddp(logging_values, next(self.net.parameters()).device)

        pbar_message_dict = get_train_loop_description_dict(logging_values,
                                                            metrics,
                                                            self.loss_logging_items_names)

        self.valid_monitored_values = sg_trainer_utils.update_monitored_values_dict(
            monitored_values_dict=self.valid_monitored_values, new_values_dict=pbar_message_dict)

        if not silent_mode and evaluation_type == EvaluationType.VALIDATION:
            progress_bar_data_loader.write("===========================================================")
            sg_trainer_utils.display_epoch_summary(epoch=context.epoch, n_digits=4,
                                                   train_monitored_values=self.train_monitored_values,
                                                   valid_monitored_values=self.valid_monitored_values)
            progress_bar_data_loader.write("===========================================================")

        return logging_values

    def _instantiate_net(self, architecture: Union[torch.nn.Module, SgModule.__class__, str], arch_params: dict,
                         checkpoint_params: dict, *args, **kwargs) -> tuple:
        """
        Instantiates nn.Module according to architecture and arch_params, and handles pretrained weights and the required
            module manipulation (i.e head replacement).

        :param architecture: String, torch.nn.Module or uninstantiated SgModule class describing the netowrks architecture.
        :param arch_params: Architecture's parameters passed to networks c'tor.
        :param checkpoint_params: checkpoint loading related parameters dictionary with 'pretrained_weights' key,
            s.t it's value is a string describing the dataset of the pretrained weights (for example "imagenent").

        :return: instantiated netowrk i.e torch.nn.Module, architecture_class (will be none when architecture is not str)

        """
        pretrained_weights = core_utils.get_param(checkpoint_params, 'pretrained_weights', default_val=None)

        if pretrained_weights is not None:
            num_classes_new_head = arch_params.num_classes
            arch_params.num_classes = PRETRAINED_NUM_CLASSES[pretrained_weights]

        if isinstance(architecture, str):
            architecture_cls = ARCHITECTURES[architecture]
            net = architecture_cls(arch_params=arch_params)
        elif isinstance(architecture, SgModule.__class__):
            net = architecture(arch_params)
        else:
            net = architecture

        if pretrained_weights:
            load_pretrained_weights(net, architecture, pretrained_weights)
            if num_classes_new_head != arch_params.num_classes:
                net.replace_head(new_num_classes=num_classes_new_head)
                arch_params.num_classes = num_classes_new_head

        return net

    def _instantiate_ema_model(self, decay: float = 0.9999, beta: float = 15, exp_activation: bool = True) -> ModelEMA:
        """Instantiate ema model for standard SgModule.
        :param decay: the maximum decay value. as the training process advances, the decay will climb towards this value
                      until the EMA_t+1 = EMA_t * decay + TRAINING_MODEL * (1- decay)
        :param beta: the exponent coefficient. The higher the beta, the sooner in the training the decay will saturate to
                     its final value. beta=15 is ~40% of the training process.
        """
        return ModelEMA(self.net, decay, beta, exp_activation)

    @property
    def get_net(self):
        """
        Getter for network.
        :return: torch.nn.Module, self.net
        """
        return self.net

    def set_net(self, net: torch.nn.Module):
        """
        Setter for network.

        :param net: torch.nn.Module, value to set net
        :return:
        """
        self.net = net

    def set_ckpt_best_name(self, ckpt_best_name):
        """
        Setter for best checkpoint filename.

        :param ckpt_best_name: str, value to set ckpt_best_name
        """
        self.ckpt_best_name = ckpt_best_name

    def set_ema(self, val: bool):
        """
        Setter for self.ema

        :param val: bool, value to set ema
        """
        self.ema = val

    def _get_context_methods(self, phase: Phase) -> ContextSgMethods:
        """
        Returns ContextSgMethods holding the methods that should be accessible through phase callbacks to the user at
         the specific phase

        :param phase: Phase, controls what methods should be returned.
        :return: ContextSgMethods holding methods from self.
        """
        if phase in [Phase.PRE_TRAINING, Phase.TRAIN_EPOCH_START, Phase.TRAIN_EPOCH_END, Phase.VALIDATION_EPOCH_END,
                     Phase.VALIDATION_EPOCH_END, Phase.POST_TRAINING, Phase.VALIDATION_END_BEST_EPOCH]:
            context_methods = ContextSgMethods(get_net=self.get_net,
                                               set_net=self.set_net,
                                               set_ckpt_best_name=self.set_ckpt_best_name,
                                               reset_best_metric=self._reset_best_metric,
                                               build_model=self.build_model,
                                               validate_epoch=self._validate_epoch,
                                               set_ema=self.set_ema)
        else:
            context_methods = ContextSgMethods()

        return context_methods

from omegaconf import DictConfig
import hydra


class Trainer:
    """
    Class for running SuperGradient's recipes.
    See train_from_recipe example in the examples directory to demonstrate it's usage.
    """

    @classmethod
    def train(cls, cfg: DictConfig) -> None:
        """
        Trains according to cfg recipe configuration.

        @param cfg: The parsed DictConfig from yaml recipe files
        @return: output of sg_model.train(...) (i.e results tuple)
        """
        # INSTANTIATE ALL OBJECTS IN CFG
        cfg = hydra.utils.instantiate(cfg)

        # CONNECT THE DATASET INTERFACE WITH DECI MODEL
        cfg.sg_model.connect_dataset_interface(cfg.dataset_interface, data_loader_num_workers=cfg.data_loader_num_workers)

        # BUILD NETWORK
        cls.build_model(cfg)

        # TRAIN
        cfg.sg_model.train(training_params=cfg.training_hyperparams)

    @classmethod
    def build_model(cls, cfg):
        cfg.sg_model.build_model(cfg.architecture, arch_params=cfg.arch_params, checkpoint_params=cfg.checkpoint_params)

@@ -44,7 +44,7 @@ class Phase(Enum):
 
 class ContextSgMethods:
     """
-    Class for delegating SgModel's methods, so that only the relevant ones are ("phase wise") are accessible.
+    Class for delegating Trainer's methods, so that only the relevant ones are ("phase wise") are accessible.
     """
     def __init__(self, **methods):
         for attr, attr_val in methods.items():

@@ -22,7 +22,7 @@ def get_ckpt_local_path(source_ckpt_folder_name: str, experiment_name: str, ckpt
         - external_checkpoint_path when external_checkpoint_path != None
 
     @param source_ckpt_folder_name: The folder where the checkpoint is saved. When set to None- uses the experiment_name.
-    @param experiment_name: experiment name attr in sg_model
+    @param experiment_name: experiment name attr in trainer
     @param ckpt_name: checkpoint filename
     @param model_checkpoints_location: S3, local ot URL
     @param external_checkpoint_path: full path to checkpoint file (that might be located outside of super_gradients/checkpoints directory)

@@ -72,8 +72,8 @@ def scaled_all_reduce(tensors: torch.Tensor, num_gpus: int):
 @torch.no_grad()
 def compute_precise_bn_stats(model: nn.Module, loader: torch.utils.data.DataLoader, precise_bn_batch_size: int, num_gpus: int):
     '''
-    :param model:                   The model being trained (ie: SgModel.net)
-    :param loader:                  Training dataloader (ie: SgModel.train_loader)
+    :param model:                   The model being trained (ie: Trainer.net)
+    :param loader:                  Training dataloader (ie: Trainer.train_loader)
     :param precise_bn_batch_size:   The effective batch size we want to calculate the batchnorm on. For example, if we are training a model
                                     on 8 gpus, with a batch of 128 on each gpu, a good rule of thumb would be to give it 8192
                                     (ie: effective_batch_size * num_gpus = batch_per_gpu * num_gpus * num_gpus).

@@ -88,7 +88,7 @@ def calibrate_model(model: torch.nn.Module, calib_data_loader: torch.utils.data.
     :param method:              str, One of [percentile, mse, entropy, max]. Statistics method for amax computation of the quantized modules
                                 (Default=percentile).
     :param num_calib_batches:   int, number of batches to collect the statistics from.
-    :param percentile:          float, percentile value to use when SgModel,quant_modules_calib_method='percentile'. Discarded when other methods are used
+    :param percentile:          float, percentile value to use when Trainer,quant_modules_calib_method='percentile'. Discarded when other methods are used
                                 (Default=99.99).
 
     """
@@ -201,7 +201,7 @@ class QATCallback(PhaseCallback):
         1. loads the best checkpoint then performs calibration.
         2. loads an external calibrated model (makes sense when start_epoch=0).
 
-    Additionally, resets SgModel's best_metric and sets ckpt_best_name to 'qat_ckpt_best.pth' so best QAT checkpoints
+    Additionally, resets Trainer's best_metric and sets ckpt_best_name to 'qat_ckpt_best.pth' so best QAT checkpoints
      will be saved separately.
 
     If performing calibration- the calibrated model is evaluated, and the metric_to_watch is logged under
@@ -225,7 +225,7 @@ class QATCallback(PhaseCallback):
 
         num_calib_batches: int, number of batches to collect the statistics from.
 
-        percentile: float, percentile value to use when SgModel,quant_modules_calib_method='percentile'.
+        percentile: float, percentile value to use when Trainer,quant_modules_calib_method='percentile'.
          Discarded when other methods are used (Default=99.99).
 
 

@@ -5,9 +5,12 @@ import time
 from dataclasses import dataclass
 from multiprocessing import Process
 from pathlib import Path
-from typing import Tuple, Union, Dict
+from typing import Tuple, Union, Dict, List, Sequence
 import random
 
+import inspect
+
+from super_gradients.common.abstractions.abstract_logger import get_logger
 from treelib import Tree
 from termcolor import colored
 import torch
@@ -16,11 +19,12 @@ from torch.utils.tensorboard import SummaryWriter
 from super_gradients.training.exceptions.dataset_exceptions import UnsupportedBatchItemsFormat
 
 
-# TODO: These utils should move to sg_model package as internal (private) helper functions
+# TODO: These utils should move to sg_trainer package as internal (private) helper functions
 
 IS_BETTER_COLOR = {True: "green", False: "red"}
 IS_GREATER_SYMBOLS = {True: "↗", False: "↘"}
 
+logger = get_logger(__name__)
 
 @dataclass
 class MonitoredValue:
@@ -329,3 +333,18 @@ def log_uncaught_exceptions(logger):
         logger.error("Uncaught exception", exc_info=(exc_type, exc_value, exc_traceback))
 
     sys.excepthook = handle_exception
+
+
+def parse_args(cfg, arg_names: Union[List[str], callable]) -> dict:
+    """
+    parse args from a config.
+    unlike get_param(), in this case only parameters that appear in the config will override default params from the function's signature
+    """
+    if not isinstance(arg_names, Sequence):
+        arg_names = list(inspect.signature(arg_names).parameters.keys())
+
+    kwargs_dict = {}
+    for arg_name in arg_names:
+        if hasattr(cfg, arg_name) and getattr(cfg, arg_name) is not None:
+            kwargs_dict[arg_name] = getattr(cfg, arg_name)
+    return kwargs_dict

@@ -14,7 +14,7 @@ from tests.unit_tests.train_with_intialized_param_args_test import TrainWithInit
 from tests.unit_tests.pretrained_models_unit_test import PretrainedModelsUnitTest
 from tests.unit_tests.lr_warmup_test import LRWarmupTest
 from tests.unit_tests.kd_ema_test import KDEMATest
-from tests.unit_tests.kd_model_test import KDModelTest
+from tests.unit_tests.kd_trainer_test import KDTrainerTest
 from tests.unit_tests.dice_loss_test import DiceLossTest
 from tests.unit_tests.iou_loss_test import IoULossTest
 from tests.unit_tests.update_param_groups_unit_test import UpdateParamGroupsTest
@@ -64,7 +64,7 @@ class CoreUnitTestSuiteRunner:
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(DiceLossTest))
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(TestViT))
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(KDEMATest))
-        self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(KDModelTest))
+        self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(KDTrainerTest))
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(TestYOLOX))
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(InitializeWithDataloadersTest))
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(LRCooldownTest))

@@ -2,18 +2,18 @@ import unittest
 
 import super_gradients
 
-from super_gradients import SgModel
+from super_gradients import Trainer
 from super_gradients.training.datasets.dataset_interfaces import LibraryDatasetInterface
 
 
 class TestCifar10Trainer(unittest.TestCase):
     def test_train_cifar10(self):
         super_gradients.init_trainer()
-        model = SgModel("test", model_checkpoints_location='local')
+        trainer = Trainer("test", model_checkpoints_location='local')
         cifar_10_dataset_interface = LibraryDatasetInterface(name="cifar10")
-        model.connect_dataset_interface(cifar_10_dataset_interface)
-        model.build_model("resnet18_cifar", arch_params={'num_classes': 10})
-        model.train(training_params={"max_epochs": 1})
+        trainer.connect_dataset_interface(cifar_10_dataset_interface)
+        trainer.build_model("resnet18_cifar", arch_params={'num_classes': 10})
+        trainer.train(training_params={"max_epochs": 1})
 
 
 if __name__ == '__main__':

@@ -5,7 +5,7 @@ import numpy as np
 from PIL import Image
 import tensorflow.keras as keras
 from super_gradients.training import MultiGPUMode
-from super_gradients.training import SgModel
+from super_gradients.training import Trainer
 from super_gradients.training.datasets.dataset_interfaces.dataset_interface import ExternalDatasetInterface, \
     ImageNetDatasetInterface
 from super_gradients.training.metrics import Accuracy, Top5
@@ -123,12 +123,12 @@ class TestExternalDatasetInterface(unittest.TestCase):
                         "greater_metric_to_watch_is_better": True}
 
         arch_params = {'num_classes': 1000}
-        model = SgModel("test", model_checkpoints_location='local',
-                        multi_gpu=MultiGPUMode.OFF)
-        model.connect_dataset_interface(dataset_interface=self.test_external_dataset_interface,
-                                        data_loader_num_workers=8)
-        model.build_model("resnet50", arch_params)
-        model.train(training_params=train_params)
+        trainer = Trainer("test", model_checkpoints_location='local',
+                          multi_gpu=MultiGPUMode.OFF)
+        trainer.connect_dataset_interface(dataset_interface=self.test_external_dataset_interface,
+                                          data_loader_num_workers=8)
+        trainer.build_model("resnet50", arch_params)
+        trainer.train(training_params=train_params)
 
 
 if __name__ == '__main__':

@@ -4,7 +4,7 @@ import unittest
 import super_gradients
 import torch
 import os
-from super_gradients import SgModel, ClassificationTestDatasetInterface
+from super_gradients import Trainer, ClassificationTestDatasetInterface
 from super_gradients.training.metrics import Accuracy, Top5
 
 
@@ -35,73 +35,65 @@ class TestTrainer(unittest.TestCase):
 
     @staticmethod
     def get_classification_trainer(name=''):
-        model = SgModel(name, model_checkpoints_location='local')
+        trainer = Trainer(name, model_checkpoints_location='local')
         dataset_params = {"batch_size": 4}
         dataset = ClassificationTestDatasetInterface(dataset_params=dataset_params)
-        model.connect_dataset_interface(dataset)
-        model.build_model("resnet18_cifar")
-        return model
+        trainer.connect_dataset_interface(dataset)
+        trainer.build_model("resnet18_cifar")
+        return trainer
 
     def test_train(self):
-        model = self.get_classification_trainer(self.folder_names[0])
-        model.train(training_params=self.training_params)
+        trainer = self.get_classification_trainer(self.folder_names[0])
+        trainer.train(training_params=self.training_params)
 
     def test_save_load(self):
-        model = self.get_classification_trainer(self.folder_names[1])
-        model.train(training_params=self.training_params)
-        model.build_model("resnet18_cifar", checkpoint_params={'load_checkpoint': True})
+        trainer = self.get_classification_trainer(self.folder_names[1])
+        trainer.train(training_params=self.training_params)
+        trainer.build_model("resnet18_cifar", checkpoint_params={'load_checkpoint': True})
 
     def test_load_only_weights_from_ckpt(self):
         # Create a checkpoint with 100% accuracy
-        model = self.get_classification_trainer(self.folder_names[2])
+        trainer = self.get_classification_trainer(self.folder_names[2])
         params = self.training_params.copy()
 
         params['max_epochs'] = 3
-        model.train(training_params=params)
+        trainer.train(training_params=params)
         # Build a model that continues the training
-        model = self.get_classification_trainer(self.folder_names[3])
-        model.build_model('resnet18_cifar', checkpoint_params={"load_checkpoint": True, "load_weights_only": False,
-                                                               "source_ckpt_folder_name": self.folder_names[2]}
-                          )
-        self.assertTrue(model.best_metric > -1)
-        self.assertTrue(model.start_epoch != 0)
+        trainer = self.get_classification_trainer(self.folder_names[3])
+        trainer.build_model('resnet18_cifar', checkpoint_params={"load_checkpoint": True, "load_weights_only": False,
+                                                                 "source_ckpt_folder_name": self.folder_names[2]}
+                            )
+        self.assertTrue(trainer.best_metric > -1)
+        self.assertTrue(trainer.start_epoch != 0)
         # start_epoch is not initialized, adding to max_epochs
         self.training_params['max_epochs'] += 3
-        model.train(training_params=self.training_params)
+        trainer.train(training_params=self.training_params)
         # Build a model that loads the weights and starts from scratch
-        model = self.get_classification_trainer(self.folder_names[4])
-        model.build_model('resnet18_cifar', checkpoint_params={"load_checkpoint": True, "load_weights_only": True,
-                                                               "source_ckpt_folder_name": self.folder_names[2]}
-                          )
-        self.assertTrue(model.best_metric == -1)
-        self.assertTrue(model.start_epoch == 0)
+        trainer = self.get_classification_trainer(self.folder_names[4])
+        trainer.build_model('resnet18_cifar', checkpoint_params={"load_checkpoint": True, "load_weights_only": True,
+                                                                 "source_ckpt_folder_name": self.folder_names[2]}
+                            )
+        self.assertTrue(trainer.best_metric == -1)
+        self.assertTrue(trainer.start_epoch == 0)
         self.training_params['max_epochs'] += 3
-        model.train(training_params=self.training_params)
+        trainer.train(training_params=self.training_params)
 
     def test_checkpoint_content(self):
         """VERIFY THAT ALL CHECKPOINTS ARE SAVED AND CONTAIN ALL THE EXPECTED KEYS"""
-        model = self.get_classification_trainer(self.folder_names[5])
+        trainer = self.get_classification_trainer(self.folder_names[5])
         params = self.training_params.copy()
         params["save_ckpt_epoch_list"] = [1]
-        model.train(training_params=params)
+        trainer.train(training_params=params)
         ckpt_filename = ['ckpt_best.pth', 'ckpt_latest.pth', 'ckpt_epoch_1.pth']
-        ckpt_paths = [os.path.join(model.checkpoints_dir_path, suf) for suf in ckpt_filename]
+        ckpt_paths = [os.path.join(trainer.checkpoints_dir_path, suf) for suf in ckpt_filename]
         for ckpt_path in ckpt_paths:
             ckpt = torch.load(ckpt_path)
             self.assertListEqual(['net', 'acc', 'epoch', 'optimizer_state_dict', 'scaler_state_dict'],
                                  list(ckpt.keys()))
-        model._save_checkpoint()
-        weights_only = torch.load(os.path.join(model.checkpoints_dir_path, 'ckpt_latest_weights_only.pth'))
+        trainer._save_checkpoint()
+        weights_only = torch.load(os.path.join(trainer.checkpoints_dir_path, 'ckpt_latest_weights_only.pth'))
         self.assertListEqual(['net'], list(weights_only.keys()))
 
-    def test_predict(self):
-        model = self.get_classification_trainer(self.folder_names[6])
-        inputs = torch.randn((5, 3, 32, 32))
-        targets = torch.randint(0, 5, (5, 1))
-        model.predict(inputs=inputs, targets=targets)
-        model.predict(inputs=inputs, targets=targets, half=True)
-        model.predict(inputs=inputs, targets=targets, half=False, verbose=True)
-
 
 if __name__ == '__main__':
     unittest.main()

@@ -3,7 +3,7 @@ from enum import Enum
 import re
 
 from super_gradients import (
-    SgModel,
+    Trainer,
     ClassificationTestDatasetInterface,
     SegmentationTestDatasetInterface,
 )
@@ -68,13 +68,13 @@ class ConversionCallbackTest(unittest.TestCase):
                 "phase_callbacks": phase_callbacks,
             }
 
-            model = SgModel(f"{architecture}_example", model_checkpoints_location="local", ckpt_root_dir=checkpoint_dir)
+            trainer = Trainer(f"{architecture}_example", model_checkpoints_location="local", ckpt_root_dir=checkpoint_dir)
             dataset = ClassificationTestDatasetInterface(dataset_params={"batch_size": 10})
 
-            model.connect_dataset_interface(dataset, data_loader_num_workers=0)
-            model.build_model(architecture=architecture, arch_params={"use_aux_heads": True, "aux_head": True})
+            trainer.connect_dataset_interface(dataset, data_loader_num_workers=0)
+            trainer.build_model(architecture=architecture, arch_params={"use_aux_heads": True, "aux_head": True})
             try:
-                model.train(train_params)
+                trainer.train(train_params)
             except Exception as e:
                 self.fail(f"Model training didn't succeed due to {e}")
             else:
@@ -103,9 +103,9 @@ class ConversionCallbackTest(unittest.TestCase):
         for architecture in SEMANTIC_SEGMENTATION:
             model_meta_data = generate_model_metadata(architecture=architecture, task=Task.SEMANTIC_SEGMENTATION)
             dataset = SegmentationTestDatasetInterface(dataset_params={"batch_size": 10})
-            model = SgModel(f"{architecture}_example", model_checkpoints_location="local", ckpt_root_dir=checkpoint_dir)
-            model.connect_dataset_interface(dataset, data_loader_num_workers=0)
-            model.build_model(architecture=architecture, arch_params={"use_aux_heads": True, "aux_head": True})
+            trainer = Trainer(f"{architecture}_example", model_checkpoints_location="local", ckpt_root_dir=checkpoint_dir)
+            trainer.connect_dataset_interface(dataset, data_loader_num_workers=0)
+            trainer.build_model(architecture=architecture, arch_params={"use_aux_heads": True, "aux_head": True})
 
             phase_callbacks = [
                 ModelConversionCheckCallback(model_meta_data=model_meta_data, opset_version=11, rtol=1, atol=1),
@@ -129,7 +129,7 @@ class ConversionCallbackTest(unittest.TestCase):
             train_params.update(custom_config)
 
             try:
-                model.train(train_params)
+                trainer.train(train_params)
             except Exception as e:
                 self.fail(f"Model training didn't succeed for {architecture} due to {e}")
             else:

@@ -1,5 +1,5 @@
 import unittest
-from super_gradients import SgModel, \
+from super_gradients import Trainer, \
     ClassificationTestDatasetInterface
 from super_gradients.training.metrics import Accuracy, Top5
 from super_gradients.training.models import ResNet18
@@ -10,12 +10,12 @@ from deci_lab_client.models import Metric, QuantizationLevel, ModelMetadata, Opt
 
 class DeciLabUploadTest(unittest.TestCase):
     def setUp(self) -> None:
-        self.model = SgModel("deci_lab_export_test_model", model_checkpoints_location='local')
+        self.trainer = Trainer("deci_lab_export_test_model", model_checkpoints_location='local')
         dataset = ClassificationTestDatasetInterface(dataset_params={"batch_size": 10})
-        self.model.connect_dataset_interface(dataset)
+        self.trainer.connect_dataset_interface(dataset)
         net = ResNet18(num_classes=5, arch_params={})
         self.optimizer = SGD(params=net.parameters(), lr=0.1)
-        self.model.build_model(net)
+        self.trainer.build_model(net)
 
     def test_train_with_deci_lab_integration(self):
         model_meta_data = ModelMetadata(name='model_for_deci_lab_upload_test',
@@ -50,7 +50,7 @@ class DeciLabUploadTest(unittest.TestCase):
                         "greater_metric_to_watch_is_better": True,
                         "phase_callbacks": [model_conversion_callback, deci_lab_callback]}
 
-        self.model.train(train_params)
+        self.trainer.train(train_params)
 
         # CLEANUP
 

@@ -1,6 +1,6 @@
 from super_gradients import ClassificationTestDatasetInterface
 from super_gradients.training import MultiGPUMode
-from super_gradients.training import SgModel
+from super_gradients.training import Trainer
 from super_gradients.training.metrics import Accuracy, Top5
 import unittest
 
@@ -23,11 +23,11 @@ class CallWrapper:
 class EMAIntegrationTest(unittest.TestCase):
 
     def _init_model(self) -> None:
-        self.model = SgModel("resnet18_cifar_ema_test", model_checkpoints_location='local',
-                             device='cpu', multi_gpu=MultiGPUMode.OFF)
+        self.trainer = Trainer("resnet18_cifar_ema_test", model_checkpoints_location='local',
+                               device='cpu', multi_gpu=MultiGPUMode.OFF)
         dataset_interface = ClassificationTestDatasetInterface({"batch_size": 32})
-        self.model.connect_dataset_interface(dataset_interface, 8)
-        self.model.build_model("resnet18_cifar")
+        self.trainer.connect_dataset_interface(dataset_interface, 8)
+        self.trainer.build_model("resnet18_cifar")
 
     @classmethod
     def tearDownClass(cls) -> None:
@@ -57,17 +57,17 @@ class EMAIntegrationTest(unittest.TestCase):
                            "greater_metric_to_watch_is_better": True}
 
         def before_test():
-            self.assertEqual(self.model.net, self.model.ema_model.ema)
+            self.assertEqual(self.trainer.net, self.trainer.ema_model.ema)
 
         def before_train_epoch():
-            self.assertNotEqual(self.model.net, self.model.ema_model.ema)
+            self.assertNotEqual(self.trainer.net, self.trainer.ema_model.ema)
 
-        self.model.test = CallWrapper(self.model.test, check_before=before_test)
-        self.model._train_epoch = CallWrapper(self.model._train_epoch, check_before=before_train_epoch)
+        self.trainer.test = CallWrapper(self.trainer.test, check_before=before_test)
+        self.trainer._train_epoch = CallWrapper(self.trainer._train_epoch, check_before=before_train_epoch)
 
-        self.model.train(training_params=training_params)
+        self.trainer.train(training_params=training_params)
 
-        self.assertIsNotNone(self.model.ema_model)
+        self.assertIsNotNone(self.trainer.ema_model)
 
 
 if __name__ == '__main__':

@@ -1,7 +1,7 @@
 import shutil
 import unittest
 import os
-from super_gradients import SgModel, ClassificationTestDatasetInterface
+from super_gradients import Trainer, ClassificationTestDatasetInterface
 from super_gradients.training.metrics import Accuracy, Top5
 
 
@@ -26,37 +26,37 @@ class LRTest(unittest.TestCase):
 
     @staticmethod
     def get_trainer(name=''):
-        model = SgModel(name, model_checkpoints_location='local')
+        trainer = Trainer(name, model_checkpoints_location='local')
         dataset_params = {"batch_size": 4}
         dataset = ClassificationTestDatasetInterface(dataset_params=dataset_params)
-        model.connect_dataset_interface(dataset)
-        model.build_model("resnet18_cifar")
-        return model
+        trainer.connect_dataset_interface(dataset)
+        trainer.build_model("resnet18_cifar")
+        return trainer
 
     def test_function_lr(self):
-        model = self.get_trainer(self.folder_name)
+        trainer = self.get_trainer(self.folder_name)
 
         def test_lr_function(initial_lr, epoch, iter, max_epoch, iters_per_epoch, **kwargs):
             return initial_lr * (1 - ((epoch * iters_per_epoch + iter) / (max_epoch * iters_per_epoch)))
 
         # test if we are able that lr_function supports functions with this structure
         training_params = {**self.training_params, "lr_mode": "function", "lr_schedule_function": test_lr_function}
-        model.train(training_params=training_params)
+        trainer.train(training_params=training_params)
 
         # test that we assert lr_function is callable
         training_params = {**self.training_params, "lr_mode": "function"}
         with self.assertRaises(AssertionError):
-            model.train(training_params=training_params)
+            trainer.train(training_params=training_params)
 
     def test_cosine_lr(self):
-        model = self.get_trainer(self.folder_name)
+        trainer = self.get_trainer(self.folder_name)
         training_params = {**self.training_params, "lr_mode": "cosine", "cosine_final_lr_ratio": 0.01}
-        model.train(training_params=training_params)
+        trainer.train(training_params=training_params)
 
     def test_step_lr(self):
-        model = self.get_trainer(self.folder_name)
+        trainer = self.get_trainer(self.folder_name)
         training_params = {**self.training_params, "lr_mode": "step", "lr_decay_factor": 0.1, "lr_updates": [4]}
-        model.train(training_params=training_params)
+        trainer.train(training_params=training_params)
 
 
 if __name__ == '__main__':

@@ -1,7 +1,7 @@
 import unittest
 import super_gradients
 from super_gradients.training import MultiGPUMode
-from super_gradients.training import SgModel
+from super_gradients.training import Trainer
 from super_gradients.training.datasets.dataset_interfaces.dataset_interface import ImageNetDatasetInterface, \
     ClassificationTestDatasetInterface, CityscapesDatasetInterface, SegmentationTestDatasetInterface, \
     CoCoSegmentationDatasetInterface, DetectionTestDatasetInterface
@@ -324,7 +324,7 @@ class PretrainedModelsTest(unittest.TestCase):
         }
 
     def test_pretrained_resnet50_imagenet(self):
-        trainer = SgModel('imagenet_pretrained_resnet50', model_checkpoints_location='local',
+        trainer = Trainer('imagenet_pretrained_resnet50', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.imagenet_dataset, data_loader_num_workers=8)
         trainer.build_model("resnet50", arch_params=self.imagenet_pretrained_arch_params["resnet"],
@@ -334,7 +334,7 @@ class PretrainedModelsTest(unittest.TestCase):
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["resnet50"], delta=0.001)
 
     def test_transfer_learning_resnet50_imagenet(self):
-        trainer = SgModel('imagenet_pretrained_resnet50_transfer_learning', model_checkpoints_location='local',
+        trainer = Trainer('imagenet_pretrained_resnet50_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
         trainer.build_model("resnet50", arch_params=self.imagenet_pretrained_arch_params["resnet"],
@@ -342,7 +342,7 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer.train(training_params=self.transfer_classification_train_params)
 
     def test_pretrained_resnet34_imagenet(self):
-        trainer = SgModel('imagenet_pretrained_resnet34', model_checkpoints_location='local',
+        trainer = Trainer('imagenet_pretrained_resnet34', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.imagenet_dataset, data_loader_num_workers=8)
         trainer.build_model("resnet34", arch_params=self.imagenet_pretrained_arch_params["resnet"],
@@ -352,7 +352,7 @@ class PretrainedModelsTest(unittest.TestCase):
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["resnet34"], delta=0.001)
 
     def test_transfer_learning_resnet34_imagenet(self):
-        trainer = SgModel('imagenet_pretrained_resnet34_transfer_learning', model_checkpoints_location='local',
+        trainer = Trainer('imagenet_pretrained_resnet34_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
         trainer.build_model("resnet34", arch_params=self.imagenet_pretrained_arch_params["resnet"],
@@ -360,7 +360,7 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer.train(training_params=self.transfer_classification_train_params)
 
     def test_pretrained_resnet18_imagenet(self):
-        trainer = SgModel('imagenet_pretrained_resnet18', model_checkpoints_location='local',
+        trainer = Trainer('imagenet_pretrained_resnet18', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.imagenet_dataset, data_loader_num_workers=8)
         trainer.build_model("resnet18", arch_params=self.imagenet_pretrained_arch_params["resnet"],
@@ -370,7 +370,7 @@ class PretrainedModelsTest(unittest.TestCase):
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["resnet18"], delta=0.001)
 
     def test_transfer_learning_resnet18_imagenet(self):
-        trainer = SgModel('imagenet_pretrained_resnet18_transfer_learning', model_checkpoints_location='local',
+        trainer = Trainer('imagenet_pretrained_resnet18_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
         trainer.build_model("resnet18", arch_params=self.imagenet_pretrained_arch_params["resnet"],
@@ -378,7 +378,7 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer.train(training_params=self.transfer_classification_train_params)
 
     def test_pretrained_regnetY800_imagenet(self):
-        trainer = SgModel('imagenet_pretrained_regnetY800', model_checkpoints_location='local',
+        trainer = Trainer('imagenet_pretrained_regnetY800', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.imagenet_dataset, data_loader_num_workers=8)
         trainer.build_model("regnetY800", arch_params=self.imagenet_pretrained_arch_params["regnet"],
@@ -388,7 +388,7 @@ class PretrainedModelsTest(unittest.TestCase):
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["regnetY800"], delta=0.001)
 
     def test_transfer_learning_regnetY800_imagenet(self):
-        trainer = SgModel('imagenet_pretrained_regnetY800_transfer_learning', model_checkpoints_location='local',
+        trainer = Trainer('imagenet_pretrained_regnetY800_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
         trainer.build_model("regnetY800", arch_params=self.imagenet_pretrained_arch_params["regnet"],
@@ -396,7 +396,7 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer.train(training_params=self.transfer_classification_train_params)
 
     def test_pretrained_regnetY600_imagenet(self):
-        trainer = SgModel('imagenet_pretrained_regnetY600', model_checkpoints_location='local',
+        trainer = Trainer('imagenet_pretrained_regnetY600', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.imagenet_dataset, data_loader_num_workers=8)
         trainer.build_model("regnetY600", arch_params=self.imagenet_pretrained_arch_params["regnet"],
@@ -406,7 +406,7 @@ class PretrainedModelsTest(unittest.TestCase):
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["regnetY600"], delta=0.001)
 
     def test_transfer_learning_regnetY600_imagenet(self):
-        trainer = SgModel('imagenet_pretrained_regnetY600_transfer_learning', model_checkpoints_location='local',
+        trainer = Trainer('imagenet_pretrained_regnetY600_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
         trainer.build_model("regnetY600", arch_params=self.imagenet_pretrained_arch_params["regnet"],
@@ -414,7 +414,7 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer.train(training_params=self.transfer_classification_train_params)
 
     def test_pretrained_regnetY400_imagenet(self):
-        trainer = SgModel('imagenet_pretrained_regnetY400', model_checkpoints_location='local',
+        trainer = Trainer('imagenet_pretrained_regnetY400', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.imagenet_dataset, data_loader_num_workers=8)
         trainer.build_model("regnetY400", arch_params=self.imagenet_pretrained_arch_params["regnet"],
@@ -424,7 +424,7 @@ class PretrainedModelsTest(unittest.TestCase):
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["regnetY400"], delta=0.001)
 
     def test_transfer_learning_regnetY400_imagenet(self):
-        trainer = SgModel('imagenet_pretrained_regnetY400_transfer_learning', model_checkpoints_location='local',
+        trainer = Trainer('imagenet_pretrained_regnetY400_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
         trainer.build_model("regnetY400", arch_params=self.imagenet_pretrained_arch_params["regnet"],
@@ -432,7 +432,7 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer.train(training_params=self.transfer_classification_train_params)
 
     def test_pretrained_regnetY200_imagenet(self):
-        trainer = SgModel('imagenet_pretrained_regnetY200', model_checkpoints_location='local',
+        trainer = Trainer('imagenet_pretrained_regnetY200', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.imagenet_dataset, data_loader_num_workers=8)
         trainer.build_model("regnetY200", arch_params=self.imagenet_pretrained_arch_params["regnet"],
@@ -442,7 +442,7 @@ class PretrainedModelsTest(unittest.TestCase):
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["regnetY200"], delta=0.001)
 
     def test_transfer_learning_regnetY200_imagenet(self):
-        trainer = SgModel('imagenet_pretrained_regnetY200_transfer_learning', model_checkpoints_location='local',
+        trainer = Trainer('imagenet_pretrained_regnetY200_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
         trainer.build_model("regnetY200", arch_params=self.imagenet_pretrained_arch_params["regnet"],
@@ -450,7 +450,7 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer.train(training_params=self.transfer_classification_train_params)
 
     def test_pretrained_repvgg_a0_imagenet(self):
-        trainer = SgModel('imagenet_pretrained_repvgg_a0', model_checkpoints_location='local',
+        trainer = Trainer('imagenet_pretrained_repvgg_a0', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.imagenet_dataset, data_loader_num_workers=8)
         trainer.build_model("repvgg_a0", arch_params=self.imagenet_pretrained_arch_params["repvgg_a0"],
@@ -460,7 +460,7 @@ class PretrainedModelsTest(unittest.TestCase):
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["repvgg_a0"], delta=0.001)
 
     def test_transfer_learning_repvgg_a0_imagenet(self):
-        trainer = SgModel('imagenet_pretrained_repvgg_a0_transfer_learning', model_checkpoints_location='local',
+        trainer = Trainer('imagenet_pretrained_repvgg_a0_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
         trainer.build_model("repvgg_a0", arch_params=self.imagenet_pretrained_arch_params["repvgg_a0"],
@@ -468,7 +468,7 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer.train(training_params=self.transfer_classification_train_params)
 
     def test_pretrained_regseg48_cityscapes(self):
-        trainer = SgModel('cityscapes_pretrained_regseg48', model_checkpoints_location='local',
+        trainer = Trainer('cityscapes_pretrained_regseg48', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.cityscapes_dataset, data_loader_num_workers=8)
         trainer.build_model("regseg48", arch_params=self.cityscapes_pretrained_arch_params["regseg48"],
@@ -479,7 +479,7 @@ class PretrainedModelsTest(unittest.TestCase):
         self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["regseg48"], delta=0.001)
 
     def test_transfer_learning_regseg48_cityscapes(self):
-        trainer = SgModel('regseg48_cityscapes_transfer_learning', model_checkpoints_location='local',
+        trainer = Trainer('regseg48_cityscapes_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.transfer_segmentation_dataset, data_loader_num_workers=8)
         trainer.build_model("regseg48", arch_params=self.cityscapes_pretrained_arch_params["regseg48"],
@@ -487,7 +487,7 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer.train(training_params=self.regseg_transfer_segmentation_train_params)
 
     def test_pretrained_ddrnet23_cityscapes(self):
-        trainer = SgModel('cityscapes_pretrained_ddrnet23', model_checkpoints_location='local',
+        trainer = Trainer('cityscapes_pretrained_ddrnet23', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.cityscapes_dataset, data_loader_num_workers=8)
         trainer.build_model("ddrnet_23", arch_params=self.cityscapes_pretrained_arch_params["ddrnet_23"],
@@ -498,7 +498,7 @@ class PretrainedModelsTest(unittest.TestCase):
         self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["ddrnet_23"], delta=0.001)
 
     def test_pretrained_ddrnet23_slim_cityscapes(self):
-        trainer = SgModel('cityscapes_pretrained_ddrnet23_slim', model_checkpoints_location='local',
+        trainer = Trainer('cityscapes_pretrained_ddrnet23_slim', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.cityscapes_dataset, data_loader_num_workers=8)
         trainer.build_model("ddrnet_23_slim", arch_params=self.cityscapes_pretrained_arch_params["ddrnet_23"],
@@ -509,7 +509,7 @@ class PretrainedModelsTest(unittest.TestCase):
         self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["ddrnet_23_slim"], delta=0.001)
 
     def test_transfer_learning_ddrnet23_cityscapes(self):
-        trainer = SgModel('cityscapes_pretrained_ddrnet23_transfer_learning', model_checkpoints_location='local',
+        trainer = Trainer('cityscapes_pretrained_ddrnet23_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.transfer_segmentation_dataset, data_loader_num_workers=8)
         trainer.build_model("ddrnet_23", arch_params=self.cityscapes_pretrained_arch_params["ddrnet_23"],
@@ -517,7 +517,7 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer.train(training_params=self.ddrnet_transfer_segmentation_train_params)
 
     def test_transfer_learning_ddrnet23_slim_cityscapes(self):
-        trainer = SgModel('cityscapes_pretrained_ddrnet23_slim_transfer_learning', model_checkpoints_location='local',
+        trainer = Trainer('cityscapes_pretrained_ddrnet23_slim_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.transfer_segmentation_dataset, data_loader_num_workers=8)
         trainer.build_model("ddrnet_23_slim", arch_params=self.cityscapes_pretrained_arch_params["ddrnet_23"],
@@ -525,7 +525,7 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer.train(training_params=self.ddrnet_transfer_segmentation_train_params)
 
     def test_pretrained_coco_segmentation_subclass_pretrained_shelfnet34_lw(self):
-        trainer = SgModel('coco_segmentation_subclass_pretrained_shelfnet34_lw', model_checkpoints_location='local',
+        trainer = Trainer('coco_segmentation_subclass_pretrained_shelfnet34_lw', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.coco_segmentation_dataset, data_loader_num_workers=8)
         trainer.build_model("shelfnet34_lw",
@@ -536,7 +536,7 @@ class PretrainedModelsTest(unittest.TestCase):
         self.assertAlmostEqual(res, self.coco_segmentation_subclass_pretrained_mious["shelfnet34_lw"], delta=0.001)
 
     def test_pretrained_efficientnet_b0_imagenet(self):
-        trainer = SgModel('imagenet_pretrained_efficientnet_b0', model_checkpoints_location='local',
+        trainer = Trainer('imagenet_pretrained_efficientnet_b0', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.imagenet_dataset, data_loader_num_workers=8)
         trainer.build_model("efficientnet_b0", arch_params=self.imagenet_pretrained_arch_params["efficientnet_b0"],
@@ -546,7 +546,7 @@ class PretrainedModelsTest(unittest.TestCase):
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["efficientnet_b0"], delta=0.001)
 
     def test_transfer_learning_efficientnet_b0_imagenet(self):
-        trainer = SgModel('imagenet_pretrained_efficientnet_b0_transfer_learning', model_checkpoints_location='local',
+        trainer = Trainer('imagenet_pretrained_efficientnet_b0_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
         trainer.build_model("efficientnet_b0", arch_params=self.imagenet_pretrained_arch_params["efficientnet_b0"],
@@ -554,7 +554,7 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer.train(training_params=self.transfer_classification_train_params)
 
     def test_pretrained_ssd_lite_mobilenet_v2_coco(self):
-        trainer = SgModel('coco_ssd_lite_mobilenet_v2', model_checkpoints_location='local',
+        trainer = Trainer('coco_ssd_lite_mobilenet_v2', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.coco_dataset['ssd_mobilenet'], data_loader_num_workers=8)
         trainer.build_model("ssd_lite_mobilenet_v2",
@@ -568,7 +568,7 @@ class PretrainedModelsTest(unittest.TestCase):
         self.assertAlmostEqual(res, self.coco_pretrained_maps["ssd_lite_mobilenet_v2"], delta=0.001)
 
     def test_transfer_learning_ssd_lite_mobilenet_v2_coco(self):
-        trainer = SgModel('coco_ssd_lite_mobilenet_v2_transfer_learning',
+        trainer = Trainer('coco_ssd_lite_mobilenet_v2_transfer_learning',
                           model_checkpoints_location='local', multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.transfer_detection_dataset,
                                           data_loader_num_workers=8)
@@ -580,7 +580,7 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer.train(training_params=self.transfer_detection_train_params['ssd_lite_mobilenet_v2'])
 
     def test_pretrained_ssd_mobilenet_v1_coco(self):
-        trainer = SgModel('coco_ssd_mobilenet_v1', model_checkpoints_location='local',
+        trainer = Trainer('coco_ssd_mobilenet_v1', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.coco_dataset['ssd_mobilenet'], data_loader_num_workers=8)
         trainer.build_model("ssd_mobilenet_v1",
@@ -595,7 +595,7 @@ class PretrainedModelsTest(unittest.TestCase):
         self.assertAlmostEqual(res, self.coco_pretrained_maps["coco_ssd_mobilenet_v1"], delta=0.001)
 
     def test_pretrained_yolox_s_coco(self):
-        trainer = SgModel('yolox_s', model_checkpoints_location='local',
+        trainer = Trainer('yolox_s', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.coco_dataset['yolox'], data_loader_num_workers=8)
         trainer.build_model("yolox_s",
@@ -607,7 +607,7 @@ class PretrainedModelsTest(unittest.TestCase):
         self.assertAlmostEqual(res, self.coco_pretrained_maps["yolox_s"], delta=0.001)
 
     def test_pretrained_yolox_m_coco(self):
-        trainer = SgModel('yolox_m', model_checkpoints_location='local',
+        trainer = Trainer('yolox_m', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.coco_dataset['yolox'], data_loader_num_workers=8)
         trainer.build_model("yolox_m",
@@ -619,7 +619,7 @@ class PretrainedModelsTest(unittest.TestCase):
         self.assertAlmostEqual(res, self.coco_pretrained_maps["yolox_m"], delta=0.001)
 
     def test_pretrained_yolox_l_coco(self):
-        trainer = SgModel('yolox_l', model_checkpoints_location='local',
+        trainer = Trainer('yolox_l', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.coco_dataset['yolox'], data_loader_num_workers=8)
         trainer.build_model("yolox_l",
@@ -631,7 +631,7 @@ class PretrainedModelsTest(unittest.TestCase):
         self.assertAlmostEqual(res, self.coco_pretrained_maps["yolox_l"], delta=0.001)
 
     def test_pretrained_yolox_n_coco(self):
-        trainer = SgModel('yolox_n', model_checkpoints_location='local',
+        trainer = Trainer('yolox_n', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.coco_dataset['yolox'], data_loader_num_workers=8)
         trainer.build_model("yolox_n",
@@ -643,7 +643,7 @@ class PretrainedModelsTest(unittest.TestCase):
         self.assertAlmostEqual(res, self.coco_pretrained_maps["yolox_n"], delta=0.001)
 
     def test_pretrained_yolox_t_coco(self):
-        trainer = SgModel('yolox_t', model_checkpoints_location='local',
+        trainer = Trainer('yolox_t', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.coco_dataset['yolox'], data_loader_num_workers=8)
         trainer.build_model("yolox_t",
@@ -655,7 +655,7 @@ class PretrainedModelsTest(unittest.TestCase):
         self.assertAlmostEqual(res, self.coco_pretrained_maps["yolox_t"], delta=0.001)
 
     def test_transfer_learning_yolox_n_coco(self):
-        trainer = SgModel('test_transfer_learning_yolox_n_coco',
+        trainer = Trainer('test_transfer_learning_yolox_n_coco',
                           model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.transfer_detection_dataset, data_loader_num_workers=8)
@@ -663,7 +663,7 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer.train(training_params=self.transfer_detection_train_params["yolox"])
 
     def test_transfer_learning_mobilenet_v3_large_imagenet(self):
-        trainer = SgModel('imagenet_pretrained_mobilenet_v3_large_transfer_learning',
+        trainer = Trainer('imagenet_pretrained_mobilenet_v3_large_transfer_learning',
                           model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
@@ -672,7 +672,7 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer.train(training_params=self.transfer_classification_train_params)
 
     def test_pretrained_mobilenet_v3_large_imagenet(self):
-        trainer = SgModel('imagenet_mobilenet_v3_large', model_checkpoints_location='local',
+        trainer = Trainer('imagenet_mobilenet_v3_large', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.imagenet_dataset, data_loader_num_workers=8)
         trainer.build_model("mobilenet_v3_large", arch_params=self.imagenet_pretrained_arch_params["mobilenet"],
@@ -682,7 +682,7 @@ class PretrainedModelsTest(unittest.TestCase):
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["mobilenet_v3_large"], delta=0.001)
 
     def test_transfer_learning_mobilenet_v3_small_imagenet(self):
-        trainer = SgModel('imagenet_pretrained_mobilenet_v3_small_transfer_learning',
+        trainer = Trainer('imagenet_pretrained_mobilenet_v3_small_transfer_learning',
                           model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
@@ -691,7 +691,7 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer.train(training_params=self.transfer_classification_train_params)
 
     def test_pretrained_mobilenet_v3_small_imagenet(self):
-        trainer = SgModel('imagenet_mobilenet_v3_small', model_checkpoints_location='local',
+        trainer = Trainer('imagenet_mobilenet_v3_small', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.imagenet_dataset, data_loader_num_workers=8)
         trainer.build_model("mobilenet_v3_small", arch_params=self.imagenet_pretrained_arch_params["mobilenet"],
@@ -701,7 +701,7 @@ class PretrainedModelsTest(unittest.TestCase):
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["mobilenet_v3_small"], delta=0.001)
 
     def test_transfer_learning_mobilenet_v2_imagenet(self):
-        trainer = SgModel('imagenet_pretrained_mobilenet_v2_transfer_learning',
+        trainer = Trainer('imagenet_pretrained_mobilenet_v2_transfer_learning',
                           model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
@@ -710,7 +710,7 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer.train(training_params=self.transfer_classification_train_params)
 
     def test_pretrained_mobilenet_v2_imagenet(self):
-        trainer = SgModel('imagenet_mobilenet_v2', model_checkpoints_location='local',
+        trainer = Trainer('imagenet_mobilenet_v2', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.imagenet_dataset, data_loader_num_workers=8)
         trainer.build_model("mobilenet_v2", arch_params=self.imagenet_pretrained_arch_params["mobilenet"],
@@ -720,7 +720,7 @@ class PretrainedModelsTest(unittest.TestCase):
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["mobilenet_v2"], delta=0.001)
 
     def test_pretrained_stdc1_seg50_cityscapes(self):
-        trainer = SgModel('cityscapes_pretrained_stdc1_seg50', model_checkpoints_location='local',
+        trainer = Trainer('cityscapes_pretrained_stdc1_seg50', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.cityscapes_dataset_rescaled50, data_loader_num_workers=8)
         trainer.build_model("stdc1_seg50", arch_params=self.cityscapes_pretrained_arch_params["stdc"],
@@ -731,7 +731,7 @@ class PretrainedModelsTest(unittest.TestCase):
         self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["stdc1_seg50"], delta=0.001)
 
     def test_transfer_learning_stdc1_seg50_cityscapes(self):
-        trainer = SgModel('cityscapes_pretrained_stdc1_seg50_transfer_learning', model_checkpoints_location='local',
+        trainer = Trainer('cityscapes_pretrained_stdc1_seg50_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.transfer_segmentation_dataset, data_loader_num_workers=8)
         trainer.build_model("stdc1_seg50", arch_params=self.cityscapes_pretrained_arch_params["stdc"],
@@ -739,7 +739,7 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer.train(training_params=self.stdc_transfer_segmentation_train_params)
 
     def test_pretrained_stdc1_seg75_cityscapes(self):
-        trainer = SgModel('cityscapes_pretrained_stdc1_seg75', model_checkpoints_location='local',
+        trainer = Trainer('cityscapes_pretrained_stdc1_seg75', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.cityscapes_dataset_rescaled75, data_loader_num_workers=8)
         trainer.build_model("stdc1_seg75", arch_params=self.cityscapes_pretrained_arch_params["stdc"],
@@ -750,7 +750,7 @@ class PretrainedModelsTest(unittest.TestCase):
         self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["stdc1_seg75"], delta=0.001)
 
     def test_transfer_learning_stdc1_seg75_cityscapes(self):
-        trainer = SgModel('cityscapes_pretrained_stdc1_seg75_transfer_learning', model_checkpoints_location='local',
+        trainer = Trainer('cityscapes_pretrained_stdc1_seg75_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.transfer_segmentation_dataset, data_loader_num_workers=8)
         trainer.build_model("stdc1_seg75", arch_params=self.cityscapes_pretrained_arch_params["stdc"],
@@ -758,7 +758,7 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer.train(training_params=self.stdc_transfer_segmentation_train_params)
 
     def test_pretrained_stdc2_seg50_cityscapes(self):
-        trainer = SgModel('cityscapes_pretrained_stdc2_seg50', model_checkpoints_location='local',
+        trainer = Trainer('cityscapes_pretrained_stdc2_seg50', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.cityscapes_dataset_rescaled50, data_loader_num_workers=8)
         trainer.build_model("stdc2_seg50", arch_params=self.cityscapes_pretrained_arch_params["stdc"],
@@ -769,7 +769,7 @@ class PretrainedModelsTest(unittest.TestCase):
         self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["stdc2_seg50"], delta=0.001)
 
     def test_transfer_learning_stdc2_seg50_cityscapes(self):
-        trainer = SgModel('cityscapes_pretrained_stdc2_seg50_transfer_learning', model_checkpoints_location='local',
+        trainer = Trainer('cityscapes_pretrained_stdc2_seg50_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.transfer_segmentation_dataset, data_loader_num_workers=8)
         trainer.build_model("stdc2_seg50", arch_params=self.cityscapes_pretrained_arch_params["stdc"],
@@ -777,7 +777,7 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer.train(training_params=self.stdc_transfer_segmentation_train_params)
 
     def test_pretrained_stdc2_seg75_cityscapes(self):
-        trainer = SgModel('cityscapes_pretrained_stdc2_seg75', model_checkpoints_location='local',
+        trainer = Trainer('cityscapes_pretrained_stdc2_seg75', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.cityscapes_dataset_rescaled75, data_loader_num_workers=8)
         trainer.build_model("stdc2_seg75", arch_params=self.cityscapes_pretrained_arch_params["stdc"],
@@ -788,7 +788,7 @@ class PretrainedModelsTest(unittest.TestCase):
         self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["stdc2_seg75"], delta=0.001)
 
     def test_transfer_learning_stdc2_seg75_cityscapes(self):
-        trainer = SgModel('cityscapes_pretrained_stdc2_seg75_transfer_learning', model_checkpoints_location='local',
+        trainer = Trainer('cityscapes_pretrained_stdc2_seg75_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.transfer_segmentation_dataset, data_loader_num_workers=8)
         trainer.build_model("stdc2_seg75", arch_params=self.cityscapes_pretrained_arch_params["stdc"],
@@ -796,7 +796,7 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer.train(training_params=self.stdc_transfer_segmentation_train_params)
 
     def test_transfer_learning_vit_base_imagenet21k(self):
-        trainer = SgModel('imagenet21k_pretrained_vit_base',
+        trainer = Trainer('imagenet21k_pretrained_vit_base',
                           model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
@@ -805,7 +805,7 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer.train(training_params=self.transfer_classification_train_params)
 
     def test_transfer_learning_vit_large_imagenet21k(self):
-        trainer = SgModel('imagenet21k_pretrained_vit_large',
+        trainer = Trainer('imagenet21k_pretrained_vit_large',
                           model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
@@ -814,7 +814,7 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer.train(training_params=self.transfer_classification_train_params)
 
     def test_pretrained_vit_base_imagenet(self):
-        trainer = SgModel('imagenet_pretrained_vit_base', model_checkpoints_location='local',
+        trainer = Trainer('imagenet_pretrained_vit_base', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.imagenet_dataset_05_mean_std, data_loader_num_workers=8)
         trainer.build_model("vit_base", arch_params=self.imagenet_pretrained_arch_params["vit_base"],
@@ -824,7 +824,7 @@ class PretrainedModelsTest(unittest.TestCase):
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["vit_base"], delta=0.001)
 
     def test_pretrained_vit_large_imagenet(self):
-        trainer = SgModel('imagenet_pretrained_vit_large', model_checkpoints_location='local',
+        trainer = Trainer('imagenet_pretrained_vit_large', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.imagenet_dataset_05_mean_std, data_loader_num_workers=8)
         trainer.build_model("vit_large", arch_params=self.imagenet_pretrained_arch_params["vit_base"],
@@ -834,7 +834,7 @@ class PretrainedModelsTest(unittest.TestCase):
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["vit_large"], delta=0.001)
 
     def test_pretrained_beit_base_imagenet(self):
-        trainer = SgModel('imagenet_pretrained_beit_base', model_checkpoints_location='local',
+        trainer = Trainer('imagenet_pretrained_beit_base', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.imagenet_dataset_05_mean_std, data_loader_num_workers=8)
         trainer.build_model("beit_base_patch16_224", arch_params=self.imagenet_pretrained_arch_params["vit_base"],
@@ -844,7 +844,7 @@ class PretrainedModelsTest(unittest.TestCase):
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["beit_base_patch16_224"], delta=0.001)
 
     def test_transfer_learning_beit_base_imagenet(self):
-        trainer = SgModel('test_transfer_learning_beit_base_imagenet',
+        trainer = Trainer('test_transfer_learning_beit_base_imagenet',
                           model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)

@@ -1,7 +1,7 @@
 import unittest
 
 from super_gradients.training.datasets.dataset_interfaces.dataset_interface import ClassificationTestDatasetInterface
-from super_gradients.training import SgModel, MultiGPUMode
+from super_gradients.training import Trainer, MultiGPUMode
 from super_gradients.training.metrics.classification_metrics import Accuracy
 import os
 from super_gradients.training.utils.quantization_utils import PostQATConversionCallback
@@ -11,12 +11,12 @@ class QATIntegrationTest(unittest.TestCase):
     def _get_trainer(self, experiment_name):
         dataset_params = {"batch_size": 10}
         dataset = ClassificationTestDatasetInterface(dataset_params=dataset_params)
-        model = SgModel(experiment_name,
-                        model_checkpoints_location='local',
-                        multi_gpu=MultiGPUMode.OFF)
-        model.connect_dataset_interface(dataset)
-        model.build_model("resnet18", checkpoint_params={"pretrained_weights": "imagenet"})
-        return model
+        trainer = Trainer(experiment_name,
+                          model_checkpoints_location='local',
+                          multi_gpu=MultiGPUMode.OFF)
+        trainer.connect_dataset_interface(dataset)
+        trainer.build_model("resnet18", checkpoint_params={"pretrained_weights": "imagenet"})
+        return trainer
 
     def _get_train_params(self, qat_params):
         train_params = {"max_epochs": 2,
@@ -38,7 +38,7 @@ class QATIntegrationTest(unittest.TestCase):
         return train_params
 
     def test_qat_from_start(self):
-        model = self._get_trainer("test_qat_from_start")
+        trainer = self._get_trainer("test_qat_from_start")
         train_params = self._get_train_params(qat_params={
             "start_epoch": 0,
             "quant_modules_calib_method": "percentile",
@@ -47,10 +47,10 @@ class QATIntegrationTest(unittest.TestCase):
             "percentile": 99.99
         })
 
-        model.train(training_params=train_params)
+        trainer.train(training_params=train_params)
 
     def test_qat_transition(self):
-        model = self._get_trainer("test_qat_transition")
+        trainer = self._get_trainer("test_qat_transition")
         train_params = self._get_train_params(qat_params={
             "start_epoch": 1,
             "quant_modules_calib_method": "percentile",
@@ -59,10 +59,10 @@ class QATIntegrationTest(unittest.TestCase):
             "percentile": 99.99
         })
 
-        model.train(training_params=train_params)
+        trainer.train(training_params=train_params)
 
     def test_qat_from_calibrated_ckpt(self):
-        model = self._get_trainer("generate_calibrated_model")
+        trainer = self._get_trainer("generate_calibrated_model")
         train_params = self._get_train_params(qat_params={
             "start_epoch": 0,
             "quant_modules_calib_method": "percentile",
@@ -71,11 +71,11 @@ class QATIntegrationTest(unittest.TestCase):
             "percentile": 99.99
         })
 
-        model.train(training_params=train_params)
+        trainer.train(training_params=train_params)
 
-        calibrated_model_path = os.path.join(model.checkpoints_dir_path, "ckpt_calibrated_percentile_99.99.pth")
+        calibrated_model_path = os.path.join(trainer.checkpoints_dir_path, "ckpt_calibrated_percentile_99.99.pth")
 
-        model = self._get_trainer("test_qat_from_calibrated_ckpt")
+        trainer = self._get_trainer("test_qat_from_calibrated_ckpt")
         train_params = self._get_train_params(qat_params={
             "start_epoch": 0,
             "quant_modules_calib_method": "percentile",
@@ -85,7 +85,7 @@ class QATIntegrationTest(unittest.TestCase):
             "percentile": 99.99
         })
 
-        model.train(training_params=train_params)
+        trainer.train(training_params=train_params)
 
 
 if __name__ == '__main__':

@@ -1,6 +1,6 @@
 import torch
 from super_gradients.training.datasets.dataset_interfaces import DatasetInterface
-from super_gradients.training.sg_model import SgModel
+from super_gradients.training.sg_trainer import Trainer
 from torchvision.models import resnet18
 import numpy as np
 
@@ -25,13 +25,13 @@ class TestDatasetInterface(DatasetInterface):
 # ------------------ Loading The Model From Model.py----------------
 arch_params = {'num_classes': 1000}
 model = resnet18()
-sg_classification_model = SgModel('Client_model_training',
-                                  model_checkpoints_location='local', device='cpu')
+trainer = Trainer('Client_model_training',
+                  model_checkpoints_location='local', device='cpu')
 # if a torch.nn.Module is provided when building the model, the model will be integrated into deci model class
-sg_classification_model.build_model(model, arch_params=arch_params)
+trainer.build_model(model, arch_params=arch_params)
 # ------------------ Loading The Dataset From Dataset.py----------------
 dataset = TestDatasetInterface()
-sg_classification_model.connect_dataset_interface(dataset)
+trainer.connect_dataset_interface(dataset)
 # ------------------ Loading The Loss From Loss.py -----------------
 loss = 'cross_entropy'
 # ------------------ Training -----------------
@@ -40,4 +40,4 @@ train_params = {"max_epochs": 100,
                 "lr_updates": [30, 60, 90, 100],
                 "lr_decay_factor": 0.1,
                 "initial_lr": 0.025, "loss": loss}
-sg_classification_model.train(train_params)
+trainer.train(train_params)

@@ -6,7 +6,7 @@ from super_gradients.training.transforms.transforms import DetectionPaddedRescal
     DetectionHSV
 from super_gradients.training.utils.detection_utils import DetectionTargetsFormat
 from super_gradients.training.utils.detection_utils import DetectionCollateFN
-from super_gradients.training.utils import sg_model_utils
+from super_gradients.training.utils import sg_trainer_utils
 from super_gradients.training import utils as core_utils
 
 
@@ -116,7 +116,7 @@ class TestDatasetInterface(unittest.TestCase):
             batch_items = next(iter(loader))
             batch_items = core_utils.tensor_container_to_device(batch_items, 'cuda', non_blocking=True)
 
-            inputs, targets, additional_batch_items = sg_model_utils.unpack_batch_items(batch_items)
+            inputs, targets, additional_batch_items = sg_trainer_utils.unpack_batch_items(batch_items)
             self.assertListEqual([batch_size, 3, image_size, image_size], list(inputs.shape))
 
     def test_pascal_voc(self):
@@ -131,7 +131,7 @@ class TestDatasetInterface(unittest.TestCase):
             batch_items = next(iter(loader))
             batch_items = core_utils.tensor_container_to_device(batch_items, 'cuda', non_blocking=True)
 
-            inputs, targets, additional_batch_items = sg_model_utils.unpack_batch_items(batch_items)
+            inputs, targets, additional_batch_items = sg_trainer_utils.unpack_batch_items(batch_items)
             self.assertListEqual([batch_size, 3, image_size, image_size], list(inputs.shape))
 
 

@@ -3,7 +3,7 @@ import unittest
 from super_gradients.training.datasets.dataset_interfaces.dataset_interface import CoCoDetectionDatasetInterface
 from super_gradients.training.metrics.detection_metrics import DetectionMetrics
 
-from super_gradients.training import SgModel
+from super_gradients.training import Trainer
 from super_gradients.training.models.detection_models.yolo_base import YoloPostPredictionCallback
 from super_gradients.training.utils.detection_utils import CrowdDetectionCollateFN, DetectionCollateFN, \
     DetectionTargetsFormat
@@ -53,11 +53,11 @@ class TestDatasetStatisticsTensorboardLogger(unittest.TestCase):
                                                                 "area_thr": 0
                                                                 })
 
-        model = SgModel('dataset_statistics_visual_test',
-                        model_checkpoints_location='local',
-                        post_prediction_callback=YoloPostPredictionCallback())
-        model.connect_dataset_interface(dataset, data_loader_num_workers=8)
-        model.build_model("yolox_s")
+        trainer = Trainer('dataset_statistics_visual_test',
+                          model_checkpoints_location='local',
+                          post_prediction_callback=YoloPostPredictionCallback())
+        trainer.connect_dataset_interface(dataset, data_loader_num_workers=8)
+        trainer.build_model("yolox_s")
 
         training_params = {"max_epochs": 1,  # we dont really need the actual training to run
                            "lr_mode": "cosine",
@@ -74,7 +74,7 @@ class TestDatasetStatisticsTensorboardLogger(unittest.TestCase):
                            "loss_logging_items_names": ["iou", "obj", "cls", "l1", "num_fg", "Loss"],
                            "metric_to_watch": "mAP@0.50:0.95",
                            }
-        model.train(training_params=training_params)
+        trainer.train(training_params=training_params)
 
 
 if __name__ == '__main__':

@@ -1,7 +1,7 @@
 import os
 import unittest
 
-from super_gradients.training import SgModel, utils as core_utils
+from super_gradients.training import Trainer, utils as core_utils
 from super_gradients.training.datasets.dataset_interfaces.dataset_interface import CoCoDetectionDatasetInterface
 from super_gradients.training.datasets.datasets_conf import COCO_DETECTION_CLASSES_LIST
 from super_gradients.training.models.detection_models.yolo_base import YoloPostPredictionCallback
@@ -46,27 +46,27 @@ class TestDetectionUtils(unittest.TestCase):
                                                                 })
 
         # Create Yolo model
-        model = SgModel('visualization_test',
-                        model_checkpoints_location='local',
-                        post_prediction_callback=YoloPostPredictionCallback())
-        model.connect_dataset_interface(dataset, data_loader_num_workers=8)
-        model.build_model("yolox_n", checkpoint_params={"pretrained_weights": "coco"})
+        trainer = Trainer('visualization_test',
+                          model_checkpoints_location='local',
+                          post_prediction_callback=YoloPostPredictionCallback())
+        trainer.connect_dataset_interface(dataset, data_loader_num_workers=8)
+        trainer.build_model("yolox_n", checkpoint_params={"pretrained_weights": "coco"})
 
         # Simulate one iteration of validation subset
-        valid_loader = model.valid_loader
+        valid_loader = trainer.valid_loader
         batch_i, (imgs, targets) = 0, next(iter(valid_loader))
-        imgs = core_utils.tensor_container_to_device(imgs, model.device)
-        targets = core_utils.tensor_container_to_device(targets, model.device)
-        output = model.net(imgs)
-        output = model.post_prediction_callback(output)
+        imgs = core_utils.tensor_container_to_device(imgs, trainer.device)
+        targets = core_utils.tensor_container_to_device(targets, trainer.device)
+        output = trainer.net(imgs)
+        output = trainer.post_prediction_callback(output)
         # Visualize the batch
         DetectionVisualization.visualize_batch(imgs, output, targets, batch_i,
-                                               COCO_DETECTION_CLASSES_LIST, model.checkpoints_dir_path)
+                                               COCO_DETECTION_CLASSES_LIST, trainer.checkpoints_dir_path)
 
         # Assert images ware created and delete them
         img_name = '{}/{}_{}.jpg'
         for i in range(4):
-            img_path = img_name.format(model.checkpoints_dir_path, batch_i, i)
+            img_path = img_name.format(trainer.checkpoints_dir_path, batch_i, i)
             self.assertTrue(os.path.exists(img_path))
             os.remove(img_path)
 

@@ -4,7 +4,7 @@ import unittest
 
 from super_gradients.training.utils.early_stopping import EarlyStop
 from super_gradients.training.utils.callbacks import Phase
-from super_gradients.training.sg_model import SgModel
+from super_gradients.training.sg_trainer import Trainer
 from super_gradients.training.datasets.dataset_interfaces import ClassificationTestDatasetInterface
 from super_gradients.training.models.classification_models.resnet import ResNet18
 from super_gradients.training.metrics import Accuracy, Top5
@@ -60,9 +60,9 @@ class EarlyStopTest(unittest.TestCase):
         Test for mode=min metric, test that training stops after no improvement in metric value for amount of `patience`
         epochs.
         """
-        model = SgModel("early_stop_test", model_checkpoints_location='local')
-        model.connect_dataset_interface(self.dataset)
-        model.build_model(self.net)
+        trainer = Trainer("early_stop_test", model_checkpoints_location='local')
+        trainer.connect_dataset_interface(self.dataset)
+        trainer.build_model(self.net)
 
         early_stop_loss = EarlyStop(Phase.VALIDATION_EPOCH_END, monitor="Loss", mode="min", patience=3, verbose=True)
         phase_callbacks = [early_stop_loss]
@@ -72,7 +72,7 @@ class EarlyStopTest(unittest.TestCase):
         train_params = self.train_params.copy()
         train_params.update({"loss": fake_loss, "phase_callbacks": phase_callbacks})
 
-        model.train(train_params)
+        trainer.train(train_params)
 
         excepted_end_epoch = 5
 
@@ -84,9 +84,9 @@ class EarlyStopTest(unittest.TestCase):
         Test for mode=max metric, test that training stops after no improvement in metric value for amount of `patience`
         epochs.
         """
-        model = SgModel("early_stop_test", model_checkpoints_location='local')
-        model.connect_dataset_interface(self.dataset)
-        model.build_model(self.net)
+        trainer = Trainer("early_stop_test", model_checkpoints_location='local')
+        trainer.connect_dataset_interface(self.dataset)
+        trainer.build_model(self.net)
 
         early_stop_acc = EarlyStop(Phase.VALIDATION_EPOCH_END, monitor="MetricTest", mode="max", patience=3,
                                    verbose=True)
@@ -98,7 +98,7 @@ class EarlyStopTest(unittest.TestCase):
         train_params.update(
             {"valid_metrics_list": [fake_metric], "metric_to_watch": "MetricTest", "phase_callbacks": phase_callbacks})
 
-        model.train(train_params)
+        trainer.train(train_params)
 
         excepted_end_epoch = 6
 
@@ -108,9 +108,9 @@ class EarlyStopTest(unittest.TestCase):
         """
         Test for mode=min metric, test that training stops after metric value reaches the `threshold` value.
         """
-        model = SgModel("early_stop_test", model_checkpoints_location='local')
-        model.connect_dataset_interface(self.dataset)
-        model.build_model(self.net)
+        trainer = Trainer("early_stop_test", model_checkpoints_location='local')
+        trainer.connect_dataset_interface(self.dataset)
+        trainer.build_model(self.net)
 
         early_stop_loss = EarlyStop(Phase.VALIDATION_EPOCH_END, monitor="Loss", mode="min", threshold=0.1, verbose=True)
         phase_callbacks = [early_stop_loss]
@@ -120,7 +120,7 @@ class EarlyStopTest(unittest.TestCase):
         train_params = self.train_params.copy()
         train_params.update({"loss": fake_loss, "phase_callbacks": phase_callbacks})
 
-        model.train(train_params)
+        trainer.train(train_params)
 
         excepted_end_epoch = 5
         # count divided by 2, because loss counter used for both train and eval.
@@ -130,9 +130,9 @@ class EarlyStopTest(unittest.TestCase):
         """
         Test for mode=max metric, test that training stops after metric value reaches the `threshold` value.
         """
-        model = SgModel("early_stop_test", model_checkpoints_location='local')
-        model.connect_dataset_interface(self.dataset)
-        model.build_model(self.net)
+        trainer = Trainer("early_stop_test", model_checkpoints_location='local')
+        trainer.connect_dataset_interface(self.dataset)
+        trainer.build_model(self.net)
 
         early_stop_acc = EarlyStop(Phase.VALIDATION_EPOCH_END, monitor="MetricTest", mode="max", threshold=0.94,
                                    verbose=True)
@@ -144,7 +144,7 @@ class EarlyStopTest(unittest.TestCase):
         train_params.update(
             {"valid_metrics_list": [fake_metric], "metric_to_watch": "MetricTest", "phase_callbacks": phase_callbacks})
 
-        model.train(train_params)
+        trainer.train(train_params)
 
         excepted_end_epoch = 7
 
@@ -155,9 +155,9 @@ class EarlyStopTest(unittest.TestCase):
         Test that training stops when monitor value is not a finite number. Test case of NaN and Inf values.
         """
         # test Nan value
-        model = SgModel("early_stop_test", model_checkpoints_location='local')
-        model.connect_dataset_interface(self.dataset)
-        model.build_model(self.net)
+        trainer = Trainer("early_stop_test", model_checkpoints_location='local')
+        trainer.connect_dataset_interface(self.dataset)
+        trainer.build_model(self.net)
 
         early_stop_loss = EarlyStop(Phase.VALIDATION_EPOCH_END, monitor="Loss", mode="min", check_finite=True,
                                     verbose=True)
@@ -168,16 +168,16 @@ class EarlyStopTest(unittest.TestCase):
         train_params = self.train_params.copy()
         train_params.update({"loss": fake_loss, "phase_callbacks": phase_callbacks})
 
-        model.train(train_params)
+        trainer.train(train_params)
 
         excepted_end_epoch = 2
 
         self.assertEqual(excepted_end_epoch, fake_loss.count // 2)
 
         # test Inf value
-        model = SgModel("early_stop_test", model_checkpoints_location='local')
-        model.connect_dataset_interface(self.dataset)
-        model.build_model(self.net)
+        trainer = Trainer("early_stop_test", model_checkpoints_location='local')
+        trainer.connect_dataset_interface(self.dataset)
+        trainer.build_model(self.net)
 
         early_stop_loss = EarlyStop(Phase.VALIDATION_EPOCH_END, monitor="Loss", mode="min", patience=3, verbose=True)
         phase_callbacks = [early_stop_loss]
@@ -187,7 +187,7 @@ class EarlyStopTest(unittest.TestCase):
         train_params = self.train_params.copy()
         train_params.update({"loss": fake_loss, "phase_callbacks": phase_callbacks})
 
-        model.train(train_params)
+        trainer.train(train_params)
 
         excepted_end_epoch = 3
         # count divided by 2, because loss counter used for both train and eval.
@@ -198,9 +198,9 @@ class EarlyStopTest(unittest.TestCase):
         Test for `min_delta` argument, metric value is considered an improvement only if
         current_value - min_delta > best_value
         """
-        model = SgModel("early_stop_test", model_checkpoints_location='local')
-        model.connect_dataset_interface(self.dataset)
-        model.build_model(self.net)
+        trainer = Trainer("early_stop_test", model_checkpoints_location='local')
+        trainer.connect_dataset_interface(self.dataset)
+        trainer.build_model(self.net)
 
         early_stop_acc = EarlyStop(Phase.VALIDATION_EPOCH_END, monitor="MetricTest", mode="max", patience=2,
                                    min_delta=0.1, verbose=True)
@@ -212,7 +212,7 @@ class EarlyStopTest(unittest.TestCase):
         train_params.update(
             {"valid_metrics_list": [fake_metric], "metric_to_watch": "MetricTest", "phase_callbacks": phase_callbacks})
 
-        model.train(train_params)
+        trainer.train(train_params)
 
         excepted_end_epoch = 5
 

@@ -2,7 +2,7 @@ import unittest
 
 import torch
 
-from super_gradients import ClassificationTestDatasetInterface, SgModel
+from super_gradients import ClassificationTestDatasetInterface, Trainer
 from super_gradients.training.metrics import Accuracy, Top5
 from super_gradients.training.models import ResNet18
 
@@ -10,13 +10,13 @@ from super_gradients.training.models import ResNet18
 class FactoriesTest(unittest.TestCase):
 
     def test_training_with_factories(self):
-        model = SgModel("test_train_with_factories", model_checkpoints_location='local')
+        trainer = Trainer("test_train_with_factories", model_checkpoints_location='local')
         dataset_params = {"batch_size": 10}
         dataset = {"classification_test_dataset": {"dataset_params": dataset_params}}
-        model.connect_dataset_interface(dataset)
+        trainer.connect_dataset_interface(dataset)
 
         net = ResNet18(num_classes=5, arch_params={})
-        model.build_model(net)
+        trainer.build_model(net)
         train_params = {"max_epochs": 2,
                         "lr_updates": [1],
                         "lr_decay_factor": 0.1,
@@ -32,12 +32,12 @@ class FactoriesTest(unittest.TestCase):
                         "loss_logging_items_names": ["Loss"], "metric_to_watch": "Accuracy",
                         "greater_metric_to_watch_is_better": True}
 
-        model.train(train_params)
+        trainer.train(train_params)
 
-        self.assertIsInstance(model.train_metrics.Accuracy, Accuracy)
-        self.assertIsInstance(model.valid_metrics.Top5, Top5)
-        self.assertIsInstance(model.dataset_interface, ClassificationTestDatasetInterface)
-        self.assertIsInstance(model.optimizer, torch.optim.ASGD)
+        self.assertIsInstance(trainer.train_metrics.Accuracy, Accuracy)
+        self.assertIsInstance(trainer.valid_metrics.Top5, Top5)
+        self.assertIsInstance(trainer.dataset_interface, ClassificationTestDatasetInterface)
+        self.assertIsInstance(trainer.optimizer, torch.optim.ASGD)
 
 
 if __name__ == '__main__':

@@ -1,5 +1,5 @@
 import unittest
-from super_gradients.training import SgModel
+from super_gradients.training import Trainer
 from super_gradients.training.metrics import Accuracy
 from super_gradients.training.datasets import ClassificationTestDatasetInterface
 from super_gradients.training.utils.callbacks import PhaseCallback, Phase, PhaseContext
@@ -34,9 +34,9 @@ class ForwardpassPrepFNTest(unittest.TestCase):
 
     def test_resizing_with_forward_pass_prep_fn(self):
         # Define Model
-        model = SgModel("ForwardpassPrepFNTest")
-        model.connect_dataset_interface(self.dataset)
-        model.build_model("resnet18", arch_params=self.arch_params)
+        trainer = Trainer("ForwardpassPrepFNTest")
+        trainer.connect_dataset_interface(self.dataset)
+        trainer.build_model("resnet18", arch_params=self.arch_params)
 
         sizes = []
         phase_callbacks = [TestInputSizesCallback(sizes)]
@@ -49,7 +49,7 @@ class ForwardpassPrepFNTest(unittest.TestCase):
                         "loss_logging_items_names": ["Loss"], "metric_to_watch": "Accuracy",
                         "greater_metric_to_watch_is_better": True, "ema": False, "phase_callbacks": phase_callbacks,
                         "pre_prediction_callback": test_forward_pass_prep_fn}
-        model.train(train_params)
+        trainer.train(train_params)
 
         # ALTHOUGH NOT SEEN IN HERE, THE 4TH EPOCH USES LR=1, SO THIS IS THE EXPECTED LIST AS WE COLLECT
         # THE LRS AFTER THE UPDATE

@@ -1,9 +1,9 @@
 import unittest
-from super_gradients import SgModel, ClassificationTestDatasetInterface
+from super_gradients import Trainer, ClassificationTestDatasetInterface
 import torch
 from torch.utils.data import TensorDataset, DataLoader
 from super_gradients.training.metrics import Accuracy
-from super_gradients.training.exceptions.sg_model_exceptions import IllegalDataloaderInitialization
+from super_gradients.training.exceptions.sg_trainer_exceptions import IllegalDataloaderInitialization
 
 
 class InitializeWithDataloadersTest(unittest.TestCase):
@@ -24,12 +24,12 @@ class InitializeWithDataloadersTest(unittest.TestCase):
         self.testcase_testloader = DataLoader(TensorDataset(inp, label))
 
     def test_interface_was_not_broken(self):
-        model = SgModel("test_interface", model_checkpoints_location='local')
+        trainer = Trainer("test_interface", model_checkpoints_location='local')
         dataset_params = {"batch_size": 10}
         dataset = ClassificationTestDatasetInterface(dataset_params=dataset_params)
-        model.connect_dataset_interface(dataset)
+        trainer.connect_dataset_interface(dataset)
 
-        model.build_model("efficientnet_b0")
+        trainer.build_model("efficientnet_b0")
         train_params = {"max_epochs": 1, "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",
@@ -37,46 +37,46 @@ class InitializeWithDataloadersTest(unittest.TestCase):
                         "train_metrics_list": [Accuracy()], "valid_metrics_list": [Accuracy()],
                         "metric_to_watch": "Accuracy",
                         "greater_metric_to_watch_is_better": True}
-        model.train(train_params)
+        trainer.train(train_params)
 
     def test_initialization_rules(self):
-        self.assertRaises(IllegalDataloaderInitialization, SgModel, "test_name", model_checkpoints_location='local',
+        self.assertRaises(IllegalDataloaderInitialization, Trainer, "test_name", model_checkpoints_location='local',
                           train_loader=self.testcase_trainloader)
-        self.assertRaises(IllegalDataloaderInitialization, SgModel, "test_name", model_checkpoints_location='local',
+        self.assertRaises(IllegalDataloaderInitialization, Trainer, "test_name", model_checkpoints_location='local',
                           valid_loader=self.testcase_validloader)
-        self.assertRaises(IllegalDataloaderInitialization, SgModel, "test_name", model_checkpoints_location='local',
+        self.assertRaises(IllegalDataloaderInitialization, Trainer, "test_name", model_checkpoints_location='local',
                           classes=self.testcase_classes)
-        self.assertRaises(IllegalDataloaderInitialization, SgModel, "test_name", model_checkpoints_location='local',
+        self.assertRaises(IllegalDataloaderInitialization, Trainer, "test_name", model_checkpoints_location='local',
                           train_loader=self.testcase_trainloader, valid_loader=self.testcase_validloader)
-        self.assertRaises(IllegalDataloaderInitialization, SgModel, "test_name", model_checkpoints_location='local',
+        self.assertRaises(IllegalDataloaderInitialization, Trainer, "test_name", model_checkpoints_location='local',
                           train_loader=self.testcase_trainloader, classes=self.testcase_classes)
-        self.assertRaises(IllegalDataloaderInitialization, SgModel, "test_name", model_checkpoints_location='local',
+        self.assertRaises(IllegalDataloaderInitialization, Trainer, "test_name", model_checkpoints_location='local',
                           valid_loader=self.testcase_validloader, classes=self.testcase_classes)
-        SgModel("test_name", model_checkpoints_location='local', train_loader=self.testcase_trainloader,
+        Trainer("test_name", model_checkpoints_location='local', train_loader=self.testcase_trainloader,
                 valid_loader=self.testcase_validloader, classes=self.testcase_classes)
-        SgModel("test_name", model_checkpoints_location='local', train_loader=self.testcase_trainloader,
+        Trainer("test_name", model_checkpoints_location='local', train_loader=self.testcase_trainloader,
                 valid_loader=self.testcase_validloader, test_loader=self.testcase_testloader,
                 classes=self.testcase_classes)
 
     def test_train_with_dataloaders(self):
-        model = SgModel(experiment_name="test_name", model_checkpoints_location="local",
-                        train_loader=self.testcase_trainloader, valid_loader=self.testcase_validloader,
-                        classes=self.testcase_classes)
+        trainer = Trainer(experiment_name="test_name", model_checkpoints_location="local",
+                          train_loader=self.testcase_trainloader, valid_loader=self.testcase_validloader,
+                          classes=self.testcase_classes)
 
-        model.build_model("resnet18")
-        model.train(training_params={"max_epochs": 2,
-                                     "lr_updates": [5, 6, 12],
-                                     "lr_decay_factor": 0.01,
-                                     "lr_mode": "step",
-                                     "initial_lr": 0.01,
-                                     "loss": "cross_entropy",
-                                     "optimizer": "SGD",
-                                     "optimizer_params": {"weight_decay": 1e-5, "momentum": 0.9},
-                                     "train_metrics_list": [Accuracy()],
-                                     "valid_metrics_list": [Accuracy()],
-                                     "metric_to_watch": "Accuracy",
-                                     "greater_metric_to_watch_is_better": True})
-        self.assertTrue(0 < model.best_metric.item() < 1)
+        trainer.build_model("resnet18")
+        trainer.train(training_params={"max_epochs": 2,
+                                       "lr_updates": [5, 6, 12],
+                                       "lr_decay_factor": 0.01,
+                                       "lr_mode": "step",
+                                       "initial_lr": 0.01,
+                                       "loss": "cross_entropy",
+                                       "optimizer": "SGD",
+                                       "optimizer_params": {"weight_decay": 1e-5, "momentum": 0.9},
+                                       "train_metrics_list": [Accuracy()],
+                                       "valid_metrics_list": [Accuracy()],
+                                       "metric_to_watch": "Accuracy",
+                                       "greater_metric_to_watch_is_better": True})
+        self.assertTrue(0 < trainer.best_metric.item() < 1)
 
 
 if __name__ == '__main__':

@@ -1,6 +1,6 @@
 import unittest
-from super_gradients.training.sg_model import SgModel
-from super_gradients.training.kd_model.kd_model import KDModel
+from super_gradients.training.sg_trainer import Trainer
+from super_gradients.training.kd_trainer.kd_trainer import KDTrainer
 import torch
 from super_gradients.training.utils.utils import check_models_have_same_weights
 from super_gradients.training.datasets.dataset_interfaces.dataset_interface import ClassificationTestDatasetInterface
@@ -11,7 +11,7 @@ from super_gradients.training.losses.kd_losses import KDLogitsLoss
 class KDEMATest(unittest.TestCase):
     @classmethod
     def setUp(cls):
-        cls.sg_trained_teacher = SgModel("sg_trained_teacher", device='cpu')
+        cls.sg_trained_teacher = Trainer("sg_trained_teacher", device='cpu')
         cls.dataset_params = {"batch_size": 5}
         cls.dataset = ClassificationTestDatasetInterface(dataset_params=cls.dataset_params)
 
@@ -29,50 +29,50 @@ class KDEMATest(unittest.TestCase):
     def test_teacher_ema_not_duplicated(self):
         """Check that the teacher EMA is a reference to the teacher net (not a copy)."""
 
-        kd_model = KDModel("test_teacher_ema_not_duplicated", device='cpu')
-        kd_model.connect_dataset_interface(self.dataset)
-        kd_model.build_model(student_architecture='resnet18',
-                             teacher_architecture='resnet50',
-                             student_arch_params={'num_classes': 1000},
-                             teacher_arch_params={'num_classes': 1000},
-                             checkpoint_params={'teacher_pretrained_weights': "imagenet"},
-                             run_teacher_on_eval=True, )
+        kd_trainer = KDTrainer("test_teacher_ema_not_duplicated", device='cpu')
+        kd_trainer.connect_dataset_interface(self.dataset)
+        kd_trainer.build_model(student_architecture='resnet18',
+                               teacher_architecture='resnet50',
+                               student_arch_params={'num_classes': 1000},
+                               teacher_arch_params={'num_classes': 1000},
+                               checkpoint_params={'teacher_pretrained_weights': "imagenet"},
+                               run_teacher_on_eval=True, )
 
-        kd_model.train(self.kd_train_params)
+        kd_trainer.train(self.kd_train_params)
 
-        self.assertTrue(kd_model.ema_model.ema.module.teacher is kd_model.net.module.teacher)
-        self.assertTrue(kd_model.ema_model.ema.module.student is not kd_model.net.module.student)
+        self.assertTrue(kd_trainer.ema_model.ema.module.teacher is kd_trainer.net.module.teacher)
+        self.assertTrue(kd_trainer.ema_model.ema.module.student is not kd_trainer.net.module.student)
 
     def test_kd_ckpt_reload_ema(self):
         """Check that the KD model load correctly from checkpoint when "load_ema_as_net=True"."""
 
         # Create a KD model and train it
-        kd_model = KDModel("test_kd_ema_ckpt_reload", device='cpu')
-        kd_model.connect_dataset_interface(self.dataset)
-        kd_model.build_model(student_architecture='resnet18',
-                             teacher_architecture='resnet50',
-                             student_arch_params={'num_classes': 1000},
-                             teacher_arch_params={'num_classes': 1000},
-                             checkpoint_params={'teacher_pretrained_weights': "imagenet"},
-                             run_teacher_on_eval=True, )
-        kd_model.train(self.kd_train_params)
-        ema_model = kd_model.ema_model.ema
-        net = kd_model.net
+        kd_trainer = KDTrainer("test_kd_ema_ckpt_reload", device='cpu')
+        kd_trainer.connect_dataset_interface(self.dataset)
+        kd_trainer.build_model(student_architecture='resnet18',
+                               teacher_architecture='resnet50',
+                               student_arch_params={'num_classes': 1000},
+                               teacher_arch_params={'num_classes': 1000},
+                               checkpoint_params={'teacher_pretrained_weights': "imagenet"},
+                               run_teacher_on_eval=True, )
+        kd_trainer.train(self.kd_train_params)
+        ema_model = kd_trainer.ema_model.ema
+        net = kd_trainer.net
 
         # Load the trained KD model
-        kd_model = KDModel("test_kd_ema_ckpt_reload", device='cpu')
-        kd_model.connect_dataset_interface(self.dataset)
-        kd_model.build_model(student_architecture='resnet18',
-                             teacher_architecture='resnet50',
-                             student_arch_params={'num_classes': 1000},
-                             teacher_arch_params={'num_classes': 1000},
-                             checkpoint_params={"load_checkpoint": True, "load_ema_as_net": True},
-                             run_teacher_on_eval=True, )
+        kd_trainer = KDTrainer("test_kd_ema_ckpt_reload", device='cpu')
+        kd_trainer.connect_dataset_interface(self.dataset)
+        kd_trainer.build_model(student_architecture='resnet18',
+                               teacher_architecture='resnet50',
+                               student_arch_params={'num_classes': 1000},
+                               teacher_arch_params={'num_classes': 1000},
+                               checkpoint_params={"load_checkpoint": True, "load_ema_as_net": True},
+                               run_teacher_on_eval=True, )
 
         # TRAIN FOR 0 EPOCHS JUST TO SEE THAT WHEN CONTINUING TRAINING EMA MODEL HAS BEEN SAVED CORRECTLY
-        kd_model.train(self.kd_train_params)
-        reloaded_ema_model = kd_model.ema_model.ema
-        reloaded_net = kd_model.net
+        kd_trainer.train(self.kd_train_params)
+        reloaded_ema_model = kd_trainer.ema_model.ema
+        reloaded_net = kd_trainer.net
 
         # trained ema == loaded ema (Should always be true as long as "ema=True" in train_params)
         self.assertTrue(check_models_have_same_weights(ema_model, reloaded_ema_model))
@@ -93,32 +93,32 @@ class KDEMATest(unittest.TestCase):
         """Check that the KD model load correctly from checkpoint when "load_ema_as_net=False"."""
 
         # Create a KD model and train it
-        kd_model = KDModel("test_kd_ema_ckpt_reload", device='cpu')
-        kd_model.connect_dataset_interface(self.dataset)
-        kd_model.build_model(student_architecture='resnet18',
-                             teacher_architecture='resnet50',
-                             student_arch_params={'num_classes': 1000},
-                             teacher_arch_params={'num_classes': 1000},
-                             checkpoint_params={'teacher_pretrained_weights': "imagenet"},
-                             run_teacher_on_eval=True, )
-        kd_model.train(self.kd_train_params)
-        ema_model = kd_model.ema_model.ema
-        net = kd_model.net
+        kd_trainer = KDTrainer("test_kd_ema_ckpt_reload", device='cpu')
+        kd_trainer.connect_dataset_interface(self.dataset)
+        kd_trainer.build_model(student_architecture='resnet18',
+                               teacher_architecture='resnet50',
+                               student_arch_params={'num_classes': 1000},
+                               teacher_arch_params={'num_classes': 1000},
+                               checkpoint_params={'teacher_pretrained_weights': "imagenet"},
+                               run_teacher_on_eval=True, )
+        kd_trainer.train(self.kd_train_params)
+        ema_model = kd_trainer.ema_model.ema
+        net = kd_trainer.net
 
         # Load the trained KD model
-        kd_model = KDModel("test_kd_ema_ckpt_reload", device='cpu')
-        kd_model.connect_dataset_interface(self.dataset)
-        kd_model.build_model(student_architecture='resnet18',
-                             teacher_architecture='resnet50',
-                             student_arch_params={'num_classes': 1000},
-                             teacher_arch_params={'num_classes': 1000},
-                             checkpoint_params={"load_checkpoint": True, "load_ema_as_net": False},
-                             run_teacher_on_eval=True, )
+        kd_trainer = KDTrainer("test_kd_ema_ckpt_reload", device='cpu')
+        kd_trainer.connect_dataset_interface(self.dataset)
+        kd_trainer.build_model(student_architecture='resnet18',
+                               teacher_architecture='resnet50',
+                               student_arch_params={'num_classes': 1000},
+                               teacher_arch_params={'num_classes': 1000},
+                               checkpoint_params={"load_checkpoint": True, "load_ema_as_net": False},
+                               run_teacher_on_eval=True, )
 
         # TRAIN FOR 0 EPOCHS JUST TO SEE THAT WHEN CONTINUING TRAINING EMA MODEL HAS BEEN SAVED CORRECTLY
-        kd_model.train(self.kd_train_params)
-        reloaded_ema_model = kd_model.ema_model.ema
-        reloaded_net = kd_model.net
+        kd_trainer.train(self.kd_train_params)
+        reloaded_ema_model = kd_trainer.ema_model.ema
+        reloaded_net = kd_trainer.net
 
         # trained ema == loaded ema (Should always be true as long as "ema=True" in train_params)
         self.assertTrue(check_models_have_same_weights(ema_model, reloaded_ema_model))

@@ -1,12 +1,12 @@
 import unittest
 import os
-from super_gradients.training.sg_model import SgModel
-from super_gradients.training.kd_model.kd_model import KDModel
+from super_gradients.training.sg_trainer import Trainer
+from super_gradients.training.kd_trainer.kd_trainer import KDTrainer
 import torch
 from super_gradients.training.utils.utils import check_models_have_same_weights
 from super_gradients.training.datasets.dataset_interfaces.dataset_interface import ClassificationTestDatasetInterface
 from super_gradients.training.metrics import Accuracy
-from super_gradients.training.exceptions.kd_model_exceptions import ArchitectureKwargsException, \
+from super_gradients.training.exceptions.kd_trainer_exceptions import ArchitectureKwargsException, \
     UnsupportedKDArchitectureException, InconsistentParamsException, UnsupportedKDModelArgException, \
     TeacherKnowledgeException
 from super_gradients.training.models.classification_models.resnet import ResNet50, ResNet18
@@ -15,10 +15,10 @@ from copy import deepcopy
 from super_gradients.training.utils.module_utils import NormalizationAdapter
 
 
-class KDModelTest(unittest.TestCase):
+class KDTrainerTest(unittest.TestCase):
     @classmethod
     def setUp(cls):
-        cls.sg_trained_teacher = SgModel("sg_trained_teacher", device='cpu')
+        cls.sg_trained_teacher = Trainer("sg_trained_teacher", device='cpu')
         cls.dataset_params = {"batch_size": 5}
         cls.dataset = ClassificationTestDatasetInterface(dataset_params=cls.dataset_params)
         cls.sg_trained_teacher.connect_dataset_interface(cls.dataset)
@@ -44,86 +44,86 @@ class KDModelTest(unittest.TestCase):
                                "greater_metric_to_watch_is_better": True, "average_best_models": False}
 
     def test_build_kd_module_with_pretrained_teacher(self):
-        kd_model = KDModel("build_kd_module_with_pretrained_teacher", device='cpu')
-        kd_model.build_model(student_architecture='resnet18', teacher_architecture='resnet50',
+        kd_trainer = KDTrainer("build_kd_module_with_pretrained_teacher", device='cpu')
+        kd_trainer.build_model(student_architecture='resnet18', teacher_architecture='resnet50',
                              student_arch_params={'num_classes': 1000}, teacher_arch_params={'num_classes': 1000},
                              checkpoint_params={'teacher_pretrained_weights': "imagenet"}
                              )
-        imagenet_resnet50_sg_model = SgModel("pretrained_resnet50")
-        imagenet_resnet50_sg_model.build_model('resnet50', arch_params={'num_classes': 1000},
+        imagenet_resnet50_trainer = Trainer("pretrained_resnet50")
+        imagenet_resnet50_trainer.build_model('resnet50', arch_params={'num_classes': 1000},
                                                checkpoint_params={'pretrained_weights': "imagenet"})
 
-        self.assertTrue(check_models_have_same_weights(kd_model.net.module.teacher,
-                                                       imagenet_resnet50_sg_model.net.module))
+        self.assertTrue(check_models_have_same_weights(kd_trainer.net.module.teacher,
+                                                       imagenet_resnet50_trainer.net.module))
 
     def test_build_kd_module_with_pretrained_student(self):
-        kd_model = KDModel("build_kd_module_with_pretrained_student", device='cpu')
-        kd_model.build_model(student_architecture='resnet18', teacher_architecture='resnet50',
+        kd_trainer = KDTrainer("build_kd_module_with_pretrained_student", device='cpu')
+        kd_trainer.build_model(student_architecture='resnet18', teacher_architecture='resnet50',
                              student_arch_params={'num_classes': 1000}, teacher_arch_params={'num_classes': 1000},
                              checkpoint_params={'student_pretrained_weights': "imagenet",
                                                 'teacher_pretrained_weights': "imagenet"}
                              )
 
-        imagenet_resnet18_sg_model = SgModel("pretrained_resnet18", device='cpu')
-        imagenet_resnet18_sg_model.build_model('resnet18', arch_params={'num_classes': 1000},
+        imagenet_resnet18_trainer = Trainer("pretrained_resnet18", device='cpu')
+        imagenet_resnet18_trainer.build_model('resnet18', arch_params={'num_classes': 1000},
                                                checkpoint_params={'pretrained_weights': "imagenet"})
 
-        self.assertTrue(check_models_have_same_weights(kd_model.net.module.student,
-                                                       imagenet_resnet18_sg_model.net.module))
+        self.assertTrue(check_models_have_same_weights(kd_trainer.net.module.student,
+                                                       imagenet_resnet18_trainer.net.module))
 
     def test_build_kd_module_pretrained_student_with_head_replacement(self):
         self.sg_trained_teacher.train(self.train_params)
         teacher_path = os.path.join(self.sg_trained_teacher.checkpoints_dir_path, 'ckpt_latest.pth')
 
-        sg_kd_model = KDModel('test_build_kd_module_student_replace_head', device='cpu')
-        sg_kd_model.build_model(student_architecture='resnet18', teacher_architecture='resnet50',
+        kd_trainer = KDTrainer('test_build_kd_module_student_replace_head', device='cpu')
+        kd_trainer.build_model(student_architecture='resnet18', teacher_architecture='resnet50',
                                 student_arch_params={'num_classes': 5}, teacher_arch_params={'num_classes': 5},
                                 checkpoint_params={'student_pretrained_weights': "imagenet",
                                                    "teacher_checkpoint_path": teacher_path}
                                 )
 
-        self.assertTrue(sg_kd_model.net.module.student.linear.out_features == 5)
+        self.assertTrue(kd_trainer.net.module.student.linear.out_features == 5)
 
     def test_build_kd_module_with_sg_trained_teacher(self):
         self.sg_trained_teacher.train(self.train_params)
         teacher_path = os.path.join(self.sg_trained_teacher.checkpoints_dir_path, 'ckpt_latest.pth')
 
-        sg_kd_model = KDModel('test_build_kd_module_with_sg_trained_teacher', device='cpu')
+        kd_trainer = KDTrainer('test_build_kd_module_with_sg_trained_teacher', device='cpu')
 
-        sg_kd_model.build_model(student_architecture='resnet18', teacher_architecture='resnet50',
+        kd_trainer.build_model(student_architecture='resnet18', teacher_architecture='resnet50',
                                 student_arch_params={'num_classes': 5}, teacher_arch_params={'num_classes': 5},
                                 checkpoint_params={"teacher_checkpoint_path": teacher_path}
                                 )
 
         self.assertTrue(
-            check_models_have_same_weights(self.sg_trained_teacher.net.module, sg_kd_model.net.module.teacher))
+            check_models_have_same_weights(self.sg_trained_teacher.net.module, kd_trainer.net.module.teacher))
 
     def test_teacher_sg_module_methods(self):
-        sg_kd_model = KDModel("test_teacher_sg_module_methods", device='cpu')
-        sg_kd_model.build_model(student_architecture='resnet18', teacher_architecture='resnet50',
+        kd_trainer = KDTrainer("test_teacher_sg_module_methods", device='cpu')
+        kd_trainer.build_model(student_architecture='resnet18', teacher_architecture='resnet50',
                                 student_arch_params={'num_classes': 1000}, teacher_arch_params={'num_classes': 1000},
                                 checkpoint_params={'teacher_pretrained_weights': "imagenet"}
                                 )
 
-        initial_param_groups = sg_kd_model.net.module.initialize_param_groups(lr=0.1, training_params={})
-        updated_param_groups = sg_kd_model.net.module.update_param_groups(param_groups=initial_param_groups, lr=0.2,
+        initial_param_groups = kd_trainer.net.module.initialize_param_groups(lr=0.1, training_params={})
+        updated_param_groups = kd_trainer.net.module.update_param_groups(param_groups=initial_param_groups, lr=0.2,
                                                                           epoch=0, iter=0, training_params={},
                                                                           total_batch=None)
 
         self.assertTrue(initial_param_groups[0]['lr'] == 0.2 == updated_param_groups[0]['lr'])
 
     def test_kd_architecture_kwarg_exception_catching(self):
-        sg_kd_model = KDModel("test_teacher_sg_module_methods", device='cpu')
+        kd_trainer = KDTrainer("test_teacher_sg_module_methods", device='cpu')
         with self.assertRaises(ArchitectureKwargsException):
-            sg_kd_model.build_model(teacher_architecture='resnet50',
+            kd_trainer.build_model(teacher_architecture='resnet50',
                                     student_arch_params={'num_classes': 5}, teacher_arch_params={'num_classes': 5},
                                     checkpoint_params={'teacher_pretrained_weights': "imagenet"}
                                     )
 
     def test_kd_unsupported_kdmodel_arg_exception_catching(self):
-        sg_kd_model = KDModel("test_teacher_sg_module_methods", device='cpu')
+        kd_trainer = KDTrainer("test_teacher_sg_module_methods", device='cpu')
         with self.assertRaises(UnsupportedKDModelArgException):
-            sg_kd_model.build_model(student_architecture='resnet18',
+            kd_trainer.build_model(student_architecture='resnet18',
                                     teacher_architecture='resnet50',
                                     student_arch_params={'num_classes': 1000},
                                     teacher_arch_params={'num_classes': 1000},
@@ -131,9 +131,9 @@ class KDModelTest(unittest.TestCase):
                                     )
 
     def test_kd_unsupported_model_exception_catching(self):
-        sg_kd_model = KDModel("test_teacher_sg_module_methods", device='cpu')
+        kd_trainer = KDTrainer("test_teacher_sg_module_methods", device='cpu')
         with self.assertRaises(UnsupportedKDArchitectureException):
-            sg_kd_model.build_model(student_architecture='resnet18',
+            kd_trainer.build_model(student_architecture='resnet18',
                                     teacher_architecture='resnet50',
                                     student_arch_params={'num_classes': 1000},
                                     teacher_arch_params={'num_classes': 1000},
@@ -142,90 +142,90 @@ class KDModelTest(unittest.TestCase):
                                     )
 
     def test_kd_inconsistent_params_exception_catching(self):
-        sg_kd_model = KDModel("test_teacher_sg_module_methods", device='cpu')
+        kd_trainer = KDTrainer("test_teacher_sg_module_methods", device='cpu')
         with self.assertRaises(InconsistentParamsException):
-            sg_kd_model.build_model(student_architecture='resnet18', teacher_architecture='resnet50',
+            kd_trainer.build_model(student_architecture='resnet18', teacher_architecture='resnet50',
                                     student_arch_params={'num_classes': 10}, teacher_arch_params={'num_classes': 1000},
                                     checkpoint_params={'teacher_pretrained_weights': "imagenet"}
                                     )
 
     def test_kd_teacher_knowledge_exception_catching(self):
-        sg_kd_model = KDModel("test_teacher_sg_module_methods", device='cpu')
+        kd_trainer = KDTrainer("test_teacher_sg_module_methods", device='cpu')
         with self.assertRaises(TeacherKnowledgeException):
-            sg_kd_model.build_model(student_architecture='resnet18', teacher_architecture='resnet50',
+            kd_trainer.build_model(student_architecture='resnet18', teacher_architecture='resnet50',
                                     student_arch_params={'num_classes': 1000}, teacher_arch_params={'num_classes': 1000}
                                     )
 
     def test_build_external_models(self):
-        sg_model = KDModel("test_training_with_external_teacher", device='cpu')
+        kd_trainer = KDTrainer("test_training_with_external_teacher", device='cpu')
         teacher_model = ResNet50(arch_params={}, num_classes=10)
         student_model = ResNet18(arch_params={}, num_classes=10)
-        sg_model.build_model(student_architecture=student_model, teacher_architecture=teacher_model,
+        kd_trainer.build_model(student_architecture=student_model, teacher_architecture=teacher_model,
                              student_arch_params={'num_classes': 10}, teacher_arch_params={'num_classes': 10}
                              )
 
         self.assertTrue(
-            check_models_have_same_weights(teacher_model, sg_model.net.module.teacher))
+            check_models_have_same_weights(teacher_model, kd_trainer.net.module.teacher))
         self.assertTrue(
-            check_models_have_same_weights(student_model, sg_model.net.module.student))
+            check_models_have_same_weights(student_model, kd_trainer.net.module.student))
 
     def test_train_kd_module_external_models(self):
-        sg_model = KDModel("test_train_kd_module_external_models", device='cpu')
+        kd_trainer = KDTrainer("test_train_kd_module_external_models", device='cpu')
         teacher_model = ResNet50(arch_params={}, num_classes=5)
         student_model = ResNet18(arch_params={}, num_classes=5)
-        sg_model.connect_dataset_interface(self.dataset)
-        sg_model.build_model(run_teacher_on_eval=True,
+        kd_trainer.connect_dataset_interface(self.dataset)
+        kd_trainer.build_model(run_teacher_on_eval=True,
                              student_arch_params={'num_classes': 5},
                              teacher_arch_params={'num_classes': 5},
                              student_architecture=deepcopy(student_model),
                              teacher_architecture=deepcopy(teacher_model),
                              )
 
-        sg_model.train(self.kd_train_params)
+        kd_trainer.train(self.kd_train_params)
 
         # TEACHER WEIGHT'S SHOULD REMAIN THE SAME
         self.assertTrue(
-            check_models_have_same_weights(teacher_model, sg_model.net.module.teacher))
+            check_models_have_same_weights(teacher_model, kd_trainer.net.module.teacher))
 
         # STUDENT WEIGHT'S SHOULD NOT REMAIN THE SAME
         self.assertFalse(
-            check_models_have_same_weights(student_model, sg_model.net.module.student))
+            check_models_have_same_weights(student_model, kd_trainer.net.module.student))
 
     def test_train_kd_module_pretrained_ckpt(self):
-        sg_model = KDModel("test_train_kd_module_pretrained_ckpt", device='cpu')
+        kd_trainer = KDTrainer("test_train_kd_module_pretrained_ckpt", device='cpu')
         teacher_model = ResNet50(arch_params={}, num_classes=5)
         teacher_path = '/tmp/teacher.pth'
         torch.save(teacher_model.state_dict(), teacher_path)
-        sg_model.connect_dataset_interface(self.dataset)
+        kd_trainer.connect_dataset_interface(self.dataset)
 
-        sg_model.build_model(student_arch_params={'num_classes': 5},
+        kd_trainer.build_model(student_arch_params={'num_classes': 5},
                              teacher_arch_params={'num_classes': 5},
                              student_architecture='resnet18',
                              teacher_architecture='resnet50',
                              checkpoint_params={"teacher_checkpoint_path": teacher_path}
                              )
-        sg_model.train(self.kd_train_params)
+        kd_trainer.train(self.kd_train_params)
 
     def test_build_model_with_input_adapter(self):
         adapter = NormalizationAdapter(mean_original=[0.485, 0.456, 0.406],
                                        std_original=[0.229, 0.224, 0.225],
                                        mean_required=[0.5, 0.5, 0.5],
                                        std_required=[0.5, 0.5, 0.5])
-        kd_model = KDModel("build_kd_module_with_with_input_adapter", device='cpu')
-        kd_model.build_model(student_architecture='resnet18', teacher_architecture='resnet50',
+        kd_trainer = KDTrainer("build_kd_module_with_with_input_adapter", device='cpu')
+        kd_trainer.build_model(student_architecture='resnet18', teacher_architecture='resnet50',
                              student_arch_params={'num_classes': 1000}, teacher_arch_params={'num_classes': 1000},
                              checkpoint_params={'teacher_pretrained_weights': "imagenet"},
                              arch_params={"teacher_input_adapter": adapter})
-        self.assertEqual(kd_model.net.module.teacher_input_adapter, adapter)
+        self.assertEqual(kd_trainer.net.module.teacher_input_adapter, adapter)
 
     def test_load_ckpt_best_for_student(self):
-        sg_model = KDModel("test_load_ckpt_best", device='cpu')
+        kd_trainer = KDTrainer("test_load_ckpt_best", device='cpu')
         teacher_model = ResNet50(arch_params={}, num_classes=5)
         teacher_path = '/tmp/teacher.pth'
         torch.save(teacher_model.state_dict(), teacher_path)
-        sg_model.connect_dataset_interface(self.dataset)
+        kd_trainer.connect_dataset_interface(self.dataset)
 
-        sg_model.build_model(student_arch_params={'num_classes': 5},
+        kd_trainer.build_model(student_arch_params={'num_classes': 5},
                              teacher_arch_params={'num_classes': 5},
                              student_architecture='resnet18',
                              teacher_architecture='resnet50',
@@ -233,24 +233,24 @@ class KDModelTest(unittest.TestCase):
                              )
         train_params = self.kd_train_params.copy()
         train_params["max_epochs"] = 1
-        sg_model.train(train_params)
-        best_student_ckpt = os.path.join(sg_model.checkpoints_dir_path, "ckpt_best.pth")
+        kd_trainer.train(train_params)
+        best_student_ckpt = os.path.join(kd_trainer.checkpoints_dir_path, "ckpt_best.pth")
 
-        student_sg_model = SgModel("studnet_sg_model")
-        student_sg_model.build_model("resnet18", arch_params={'num_classes': 5},
+        student_trainer = Trainer("studnet_trainer")
+        student_trainer.build_model("resnet18", arch_params={'num_classes': 5},
                                      checkpoint_params={"load_checkpoint": True, "external_checkpoint_path": best_student_ckpt})
 
         self.assertTrue(
-            check_models_have_same_weights(student_sg_model.net.module, sg_model.net.module.student))
+            check_models_have_same_weights(student_trainer.net.module, kd_trainer.net.module.student))
 
     def test_load_ckpt_best_for_student_with_ema(self):
-        sg_model = KDModel("test_load_ckpt_best_for_student_with_ema", device='cpu')
+        kd_trainer = KDTrainer("test_load_ckpt_best_for_student_with_ema", device='cpu')
         teacher_model = ResNet50(arch_params={}, num_classes=5)
         teacher_path = '/tmp/teacher.pth'
         torch.save(teacher_model.state_dict(), teacher_path)
-        sg_model.connect_dataset_interface(self.dataset)
+        kd_trainer.connect_dataset_interface(self.dataset)
 
-        sg_model.build_model(student_arch_params={'num_classes': 5},
+        kd_trainer.build_model(student_arch_params={'num_classes': 5},
                              teacher_arch_params={'num_classes': 5},
                              student_architecture='resnet18',
                              teacher_architecture='resnet50',
@@ -259,14 +259,14 @@ class KDModelTest(unittest.TestCase):
         train_params = self.kd_train_params.copy()
         train_params["max_epochs"] = 1
         train_params["ema"] = True
-        sg_model.train(train_params)
-        best_student_ckpt = os.path.join(sg_model.checkpoints_dir_path, "ckpt_best.pth")
+        kd_trainer.train(train_params)
+        best_student_ckpt = os.path.join(kd_trainer.checkpoints_dir_path, "ckpt_best.pth")
 
-        student_sg_model = SgModel("studnet_sg_model")
-        student_sg_model.build_model("resnet18", arch_params={'num_classes': 5},
+        student_trainer = Trainer("studnet_trainer")
+        student_trainer.build_model("resnet18", arch_params={'num_classes': 5},
                                      checkpoint_params={"load_checkpoint": True, "external_checkpoint_path": best_student_ckpt})
         self.assertTrue(
-            check_models_have_same_weights(student_sg_model.net.module, sg_model.ema_model.ema.module.student))
+            check_models_have_same_weights(student_trainer.net.module, kd_trainer.ema_model.ema.module.student))
 
 
 if __name__ == '__main__':

@@ -2,11 +2,11 @@ import shutil
 import tempfile
 import unittest
 import os
-from super_gradients.training import SgModel
+from super_gradients.training import Trainer
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from super_gradients.training.sg_model.sg_model import StrictLoad
+from super_gradients.training.sg_trainer.sg_trainer import StrictLoad
 
 
 class Net(nn.Module):
@@ -55,15 +55,15 @@ class LoadCheckpointFromDirectPathTest(unittest.TestCase):
         # Make sure we initialized a model with different weights
         assert not self.check_models_have_same_weights(new_torch_net, self.original_torch_net)
 
-        # Build the SgModel and load the checkpoint
-        model = SgModel("load_checkpoint_test", model_checkpoints_location='local')
-        model.build_model(new_torch_net, arch_params={'num_classes': 10},
-                          checkpoint_params={'external_checkpoint_path': self.checkpoint_path,
-                                             'load_checkpoint': True,
-                                             'strict_load': StrictLoad.NO_KEY_MATCHING})
+        # Build the Trainer and load the checkpoint
+        trainer = Trainer("load_checkpoint_test", model_checkpoints_location='local')
+        trainer.build_model(new_torch_net, arch_params={'num_classes': 10},
+                            checkpoint_params={'external_checkpoint_path': self.checkpoint_path,
+                                               'load_checkpoint': True,
+                                               'strict_load': StrictLoad.NO_KEY_MATCHING})
 
         # Assert the weights were loaded correctly
-        assert self.check_models_have_same_weights(model.net, self.original_torch_net)
+        assert self.check_models_have_same_weights(trainer.net, self.original_torch_net)
 
     def check_models_have_same_weights(self, model_1, model_2):
         model_1, model_2 = model_1.to('cpu'), model_2.to('cpu')

@@ -1,5 +1,5 @@
 import unittest
-from super_gradients.training import SgModel
+from super_gradients.training import Trainer
 from super_gradients.training.metrics import Accuracy, Top5
 from super_gradients.training.utils.utils import check_models_have_same_weights
 from super_gradients.training.datasets import ClassificationTestDatasetInterface
@@ -20,24 +20,24 @@ class LoadCheckpointWithEmaTest(unittest.TestCase):
     def test_ema_ckpt_reload(self):
         # Define Model
         net = LeNet()
-        model = SgModel("ema_ckpt_test", model_checkpoints_location='local')
+        trainer = Trainer("ema_ckpt_test", model_checkpoints_location='local')
 
-        model.connect_dataset_interface(self.dataset)
-        model.build_model(net, arch_params={'num_classes': 10})
+        trainer.connect_dataset_interface(self.dataset)
+        trainer.build_model(net, arch_params={'num_classes': 10})
 
-        model.train(self.train_params)
+        trainer.train(self.train_params)
 
-        ema_model = model.ema_model.ema
+        ema_model = trainer.ema_model.ema
 
         net = LeNet()
-        model = SgModel("ema_ckpt_test", model_checkpoints_location='local')
-        model.build_model(net, arch_params={'num_classes': 10}, checkpoint_params={'load_checkpoint': True})
-        model.connect_dataset_interface(self.dataset)
+        trainer = Trainer("ema_ckpt_test", model_checkpoints_location='local')
+        trainer.build_model(net, arch_params={'num_classes': 10}, checkpoint_params={'load_checkpoint': True})
+        trainer.connect_dataset_interface(self.dataset)
 
         # TRAIN FOR 0 EPOCHS JUST TO SEE THAT WHEN CONTINUING TRAINING EMA MODEL HAS BEEN SAVED CORRECTLY
-        model.train(self.train_params)
+        trainer.train(self.train_params)
 
-        reloaded_ema_model = model.ema_model.ema
+        reloaded_ema_model = trainer.ema_model.ema
 
         # ASSERT RELOADED EMA MODEL HAS THE SAME WEIGHTS AS THE EMA MODEL SAVED IN FIRST PART OF TRAINING
         assert check_models_have_same_weights(ema_model, reloaded_ema_model)

@@ -1,5 +1,5 @@
 import unittest
-from super_gradients.training import SgModel
+from super_gradients.training import Trainer
 from super_gradients.training.metrics import Accuracy
 from super_gradients.training.datasets import ClassificationTestDatasetInterface
 from super_gradients.training.models import LeNet
@@ -15,9 +15,9 @@ class LRCooldownTest(unittest.TestCase):
     def test_lr_cooldown_with_lr_scheduling(self):
         # Define Model
         net = LeNet()
-        model = SgModel("lr_warmup_test", model_checkpoints_location='local')
-        model.connect_dataset_interface(self.dataset)
-        model.build_model(net, arch_params=self.arch_params)
+        trainer = Trainer("lr_warmup_test", model_checkpoints_location='local')
+        trainer.connect_dataset_interface(self.dataset)
+        trainer.build_model(net, arch_params=self.arch_params)
 
         lrs = []
         phase_callbacks = [TestLRCallback(lr_placeholder=lrs)]
@@ -31,7 +31,7 @@ class LRCooldownTest(unittest.TestCase):
                         "greater_metric_to_watch_is_better": True, "ema": False, "phase_callbacks": phase_callbacks}
 
         expected_lrs = [0.25, 0.5, 0.75, 0.9236067977499791, 0.4763932022500211, 0.4763932022500211, 0.4763932022500211]
-        model.train(train_params)
+        trainer.train(train_params)
 
         # ALTHOUGH NOT SEEN IN HERE, THE 4TH EPOCH USES LR=1, SO THIS IS THE EXPECTED LIST AS WE COLLECT
         # THE LRS AFTER THE UPDATE

@@ -2,7 +2,7 @@ import unittest
 
 import numpy as np
 
-from super_gradients.training import SgModel
+from super_gradients.training import Trainer
 from super_gradients.training.metrics import Accuracy
 from super_gradients.training.datasets import ClassificationTestDatasetInterface
 from super_gradients.training.models import LeNet
@@ -42,9 +42,9 @@ class LRWarmupTest(unittest.TestCase):
     def test_lr_warmup(self):
         # Define Model
         net = LeNet()
-        model = SgModel("lr_warmup_test", model_checkpoints_location='local')
-        model.connect_dataset_interface(self.dataset)
-        model.build_model(net, arch_params=self.arch_params)
+        trainer = Trainer("lr_warmup_test", model_checkpoints_location='local')
+        trainer.connect_dataset_interface(self.dataset)
+        trainer.build_model(net, arch_params=self.arch_params)
 
         lrs = []
         phase_callbacks = [TestLRCallback(lr_placeholder=lrs)]
@@ -58,15 +58,15 @@ class LRWarmupTest(unittest.TestCase):
                         "warmup_mode": "linear_step"}
 
         expected_lrs = [0.25, 0.5, 0.75, 1.0, 1.0]
-        model.train(train_params)
+        trainer.train(train_params)
         self.assertListEqual(lrs, expected_lrs)
 
     def test_lr_warmup_with_lr_scheduling(self):
         # Define Model
         net = LeNet()
-        model = SgModel("lr_warmup_test", model_checkpoints_location='local')
-        model.connect_dataset_interface(self.dataset)
-        model.build_model(net, arch_params=self.arch_params)
+        trainer = Trainer("lr_warmup_test", model_checkpoints_location='local')
+        trainer.connect_dataset_interface(self.dataset)
+        trainer.build_model(net, arch_params=self.arch_params)
 
         lrs = []
         phase_callbacks = [TestLRCallback(lr_placeholder=lrs)]
@@ -80,7 +80,7 @@ class LRWarmupTest(unittest.TestCase):
                         "warmup_mode": "linear_step"}
 
         expected_lrs = [0.25, 0.5, 0.75, 0.9236067977499791, 0.4763932022500211]
-        model.train(train_params)
+        trainer.train(train_params)
 
         # ALTHOUGH NOT SEEN IN HERE, THE 4TH EPOCH USES LR=1, SO THIS IS THE EXPECTED LIST AS WE COLLECT
         # THE LRS AFTER THE UPDATE
@@ -89,9 +89,9 @@ class LRWarmupTest(unittest.TestCase):
     def test_warmup_initial_lr(self):
         # Define Model
         net = LeNet()
-        model = SgModel("test_warmup_initial_lr", model_checkpoints_location='local')
-        model.connect_dataset_interface(self.dataset)
-        model.build_model(net, arch_params=self.arch_params)
+        trainer = Trainer("test_warmup_initial_lr", model_checkpoints_location='local')
+        trainer.connect_dataset_interface(self.dataset)
+        trainer.build_model(net, arch_params=self.arch_params)
 
         lrs = []
         phase_callbacks = [TestLRCallback(lr_placeholder=lrs)]
@@ -105,15 +105,15 @@ class LRWarmupTest(unittest.TestCase):
                         "warmup_mode": "linear_step", "initial_lr": 1, "warmup_initial_lr": 4.}
 
         expected_lrs = [4., 3., 2., 1., 1.]
-        model.train(train_params)
+        trainer.train(train_params)
         self.assertListEqual(lrs, expected_lrs)
 
     def test_custom_lr_warmup(self):
         # Define Model
         net = LeNet()
-        model = SgModel("custom_lr_warmup_test", model_checkpoints_location='local')
-        model.connect_dataset_interface(self.dataset)
-        model.build_model(net, arch_params=self.arch_params)
+        trainer = Trainer("custom_lr_warmup_test", model_checkpoints_location='local')
+        trainer.connect_dataset_interface(self.dataset)
+        trainer.build_model(net, arch_params=self.arch_params)
 
         lrs = []
         phase_callbacks = [TestLRCallback(lr_placeholder=lrs)]
@@ -127,7 +127,7 @@ class LRWarmupTest(unittest.TestCase):
                         "warmup_mode": ExponentialWarmupLRCallback, "initial_lr": 1., "warmup_initial_lr": 0.1}
 
         expected_lrs = [0.1, 0.18102751585334242, 0.40128313980266034, 1.0, 1.0]
-        model.train(train_params)
+        trainer.train(train_params)
         self.assertListEqual(lrs, expected_lrs)
 
 

@@ -1,6 +1,6 @@
 import unittest
 from super_gradients.training.utils.callbacks import PhaseContextTestCallback, Phase
-from super_gradients import SgModel, \
+from super_gradients import Trainer, \
     ClassificationTestDatasetInterface
 from super_gradients.training.metrics import Accuracy, Top5
 from super_gradients.training.models import ResNet18
@@ -11,13 +11,13 @@ from torchmetrics import MetricCollection
 
 class PhaseContextTest(unittest.TestCase):
     def context_information_in_train_test(self):
-        model = SgModel("context_information_in_train_test", model_checkpoints_location='local')
+        trainer = Trainer("context_information_in_train_test", model_checkpoints_location='local')
         dataset_params = {"batch_size": 10}
         dataset = ClassificationTestDatasetInterface(dataset_params=dataset_params)
-        model.connect_dataset_interface(dataset)
+        trainer.connect_dataset_interface(dataset)
 
         net = ResNet18(num_classes=5, arch_params={})
-        model.build_model(net)
+        trainer.build_model(net)
 
         phase_callbacks = [PhaseContextTestCallback(Phase.TRAIN_BATCH_END),
                            PhaseContextTestCallback(Phase.TRAIN_BATCH_STEP),
@@ -32,8 +32,8 @@ class PhaseContextTest(unittest.TestCase):
                         "loss_logging_items_names": ["Loss"], "metric_to_watch": "Top5",
                         "greater_metric_to_watch_is_better": True, "phase_callbacks": phase_callbacks}
 
-        model.train(train_params)
-        context_callbacks = list(filter(lambda cb: isinstance(cb, PhaseContextTestCallback), model.phase_callbacks))
+        trainer.train(train_params)
+        context_callbacks = list(filter(lambda cb: isinstance(cb, PhaseContextTestCallback), trainer.phase_callbacks))
 
         # CHECK THAT PHASE CONTEXES HAVE THE EXACT INFORMATION THERY'RE SUPPOSE TO HOLD
         for phase_callback in context_callbacks:

@@ -1,5 +1,5 @@
 import unittest
-from super_gradients.training import SgModel
+from super_gradients.training import Trainer
 from super_gradients.training.metrics import Accuracy
 from super_gradients.training.datasets import ClassificationTestDatasetInterface
 from super_gradients.training.models import LeNet
@@ -8,7 +8,7 @@ from super_gradients.training.utils.callbacks import Phase, PhaseCallback, Phase
 
 class ContextMethodsCheckerCallback(PhaseCallback):
     """
-    Callback for checking that at a certain phase specific SgModel methods are accessible.
+    Callback for checking that at a certain phase specific Trainer methods are accessible.
     """
 
     def __init__(self, phase: Phase, accessible_method_names: list, non_accessible_method_names: list):
@@ -35,9 +35,9 @@ class ContextMethodsTest(unittest.TestCase):
 
     def test_access_to_methods_by_phase(self):
         net = LeNet()
-        model = SgModel("test_access_to_methods_by_phase", model_checkpoints_location='local')
-        model.connect_dataset_interface(self.dataset)
-        model.build_model(net, arch_params=self.arch_params)
+        trainer = Trainer("test_access_to_methods_by_phase", model_checkpoints_location='local')
+        trainer.connect_dataset_interface(self.dataset)
+        trainer.build_model(net, arch_params=self.arch_params)
 
         phase_callbacks = []
         for phase in Phase:
@@ -68,7 +68,7 @@ class ContextMethodsTest(unittest.TestCase):
                         "loss_logging_items_names": ["Loss"], "metric_to_watch": "Accuracy",
                         "greater_metric_to_watch_is_better": True, "ema": False, "phase_callbacks": phase_callbacks}
 
-        model.train(train_params)
+        trainer.train(train_params)
         for phase_callback in phase_callbacks:
             if isinstance(phase_callback, ContextMethodsCheckerCallback):
                 self.assertTrue(phase_callback.result)

@@ -1,7 +1,7 @@
 import unittest
 import super_gradients
 from super_gradients.training import MultiGPUMode
-from super_gradients.training import SgModel
+from super_gradients.training import Trainer
 from super_gradients.training.datasets.dataset_interfaces.dataset_interface import ClassificationTestDatasetInterface
 from super_gradients.training.metrics import Accuracy
 import os
@@ -16,7 +16,7 @@ class PretrainedModelsUnitTest(unittest.TestCase):
         self.test_dataset = ClassificationTestDatasetInterface(classes=range(1000))
 
     def test_pretrained_resnet50_imagenet(self):
-        trainer = SgModel('imagenet_pretrained_resnet50_unit_test', model_checkpoints_location='local',
+        trainer = Trainer('imagenet_pretrained_resnet50_unit_test', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.test_dataset, data_loader_num_workers=8)
         trainer.build_model("resnet50", checkpoint_params={"pretrained_weights": "imagenet"})
@@ -24,7 +24,7 @@ class PretrainedModelsUnitTest(unittest.TestCase):
                      metrics_progress_verbose=True)
 
     def test_pretrained_regnetY800_imagenet(self):
-        trainer = SgModel('imagenet_pretrained_regnetY800_unit_test', model_checkpoints_location='local',
+        trainer = Trainer('imagenet_pretrained_regnetY800_unit_test', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.test_dataset, data_loader_num_workers=8)
         trainer.build_model("regnetY800", checkpoint_params={"pretrained_weights": "imagenet"})
@@ -32,7 +32,7 @@ class PretrainedModelsUnitTest(unittest.TestCase):
                      metrics_progress_verbose=True)
 
     def test_pretrained_repvgg_a0_imagenet(self):
-        trainer = SgModel('imagenet_pretrained_repvgg_a0_unit_test', model_checkpoints_location='local',
+        trainer = Trainer('imagenet_pretrained_repvgg_a0_unit_test', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.test_dataset, data_loader_num_workers=8)
         trainer.build_model("repvgg_a0", checkpoint_params={"pretrained_weights": "imagenet"},

@@ -1,6 +1,6 @@
 import unittest
 import os
-from super_gradients.training import SgModel
+from super_gradients.training import Trainer
 from super_gradients.training.datasets.dataset_interfaces.dataset_interface import ClassificationTestDatasetInterface
 from super_gradients.training.metrics import Accuracy, Top5
 
@@ -18,19 +18,19 @@ class SaveCkptListUnitTest(unittest.TestCase):
                         "greater_metric_to_watch_is_better": True}
 
         # Define Model
-        model = SgModel("save_ckpt_test", model_checkpoints_location='local')
+        trainer = Trainer("save_ckpt_test", model_checkpoints_location='local')
 
         # Connect Dataset
         dataset = ClassificationTestDatasetInterface()
-        model.connect_dataset_interface(dataset, data_loader_num_workers=8)
+        trainer.connect_dataset_interface(dataset, data_loader_num_workers=8)
 
         # Build Model
-        model.build_model("resnet18_cifar")
+        trainer.build_model("resnet18_cifar")
 
         # Train Model (and save ckpt_epoch_list)
-        model.train(training_params=train_params)
+        trainer.train(training_params=train_params)
 
-        dir_path = model.checkpoints_dir_path
+        dir_path = trainer.checkpoints_dir_path
         self.file_names_list = [dir_path + f'/ckpt_epoch_{epoch}.pth' for epoch in train_params["save_ckpt_epoch_list"]]
 
     def test_save_ckpt_epoch_list(self):

@@ -4,12 +4,12 @@ import unittest
 import os
 
 from super_gradients.common.sg_loggers import BaseSGLogger
-from super_gradients.training import SgModel
+from super_gradients.training import Trainer
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from super_gradients.training.sg_model.sg_model import StrictLoad
+from super_gradients.training.sg_trainer.sg_trainer import StrictLoad
 from super_gradients.training.utils import HpmStruct
 
 
@@ -51,14 +51,14 @@ class StrictLoadEnumTest(unittest.TestCase):
         # Save the model's state_dict checkpoint with different keys
         torch.save(cls.change_state_dict_keys(cls.original_torch_net.state_dict()), cls.checkpoint_diff_keys_path)
 
-        # Save the model's state_dict checkpoint in SgModel format
-        cls.sg_model = SgModel("load_checkpoint_test", model_checkpoints_location='local')  # Saves in /checkpoints
-        cls.sg_model.build_model(cls.original_torch_net, arch_params={'num_classes': 10})
+        # Save the model's state_dict checkpoint in Trainer format
+        cls.trainer = Trainer("load_checkpoint_test", model_checkpoints_location='local')  # Saves in /checkpoints
+        cls.trainer.build_model(cls.original_torch_net, arch_params={'num_classes': 10})
         # FIXME: after uniting init and build_model we should remove this
-        cls.sg_model.sg_logger = BaseSGLogger('project_name', 'load_checkpoint_test', 'local', resumed=False,
-                                              training_params=HpmStruct(max_epochs=10),
-                                              checkpoints_dir_path=cls.sg_model.checkpoints_dir_path)
-        cls.sg_model._save_checkpoint()
+        cls.trainer.sg_logger = BaseSGLogger('project_name', 'load_checkpoint_test', 'local', resumed=False,
+                                             training_params=HpmStruct(max_epochs=10),
+                                             checkpoints_dir_path=cls.trainer.checkpoints_dir_path)
+        cls.trainer._save_checkpoint()
 
     @classmethod
     def tearDownClass(cls):
@@ -96,15 +96,15 @@ class StrictLoadEnumTest(unittest.TestCase):
         # Make sure we initialized a model with different weights
         assert not self.check_models_have_same_weights(new_torch_net, self.original_torch_net)
 
-        # Build the SgModel and load the checkpoint
-        model = SgModel(self.experiment_name, model_checkpoints_location='local',
-                        ckpt_name='ckpt_latest_weights_only.pth')
-        model.build_model(new_torch_net, arch_params={'num_classes': 10},
-                          checkpoint_params={'strict_load': StrictLoad.ON,
-                                             'load_checkpoint': True})
+        # Build the Trainer and load the checkpoint
+        trainer = Trainer(self.experiment_name, model_checkpoints_location='local',
+                          ckpt_name='ckpt_latest_weights_only.pth')
+        trainer.build_model(new_torch_net, arch_params={'num_classes': 10},
+                            checkpoint_params={'strict_load': StrictLoad.ON,
+                                               'load_checkpoint': True})
 
         # Assert the weights were loaded correctly
-        assert self.check_models_have_same_weights(model.net, self.original_torch_net)
+        assert self.check_models_have_same_weights(trainer.net, self.original_torch_net)
 
     def test_strict_load_off(self):
         # Define Model
@@ -113,15 +113,15 @@ class StrictLoadEnumTest(unittest.TestCase):
         # Make sure we initialized a model with different weights
         assert not self.check_models_have_same_weights(new_torch_net, self.original_torch_net)
 
-        # Build the SgModel and load the checkpoint
-        model = SgModel(self.experiment_name, model_checkpoints_location='local',
-                        ckpt_name='ckpt_latest_weights_only.pth')
-        model.build_model(new_torch_net, arch_params={'num_classes': 10},
-                          checkpoint_params={'strict_load': StrictLoad.OFF,
-                                             'load_checkpoint': True})
+        # Build the Trainer and load the checkpoint
+        trainer = Trainer(self.experiment_name, model_checkpoints_location='local',
+                          ckpt_name='ckpt_latest_weights_only.pth')
+        trainer.build_model(new_torch_net, arch_params={'num_classes': 10},
+                            checkpoint_params={'strict_load': StrictLoad.OFF,
+                                               'load_checkpoint': True})
 
         # Assert the weights were loaded correctly
-        assert self.check_models_have_same_weights(model.net, self.original_torch_net)
+        assert self.check_models_have_same_weights(trainer.net, self.original_torch_net)
 
     def test_strict_load_no_key_matching_external_checkpoint(self):
         # Define Model
@@ -130,15 +130,15 @@ class StrictLoadEnumTest(unittest.TestCase):
         # Make sure we initialized a model with different weights
         assert not self.check_models_have_same_weights(new_torch_net, self.original_torch_net)
 
-        # Build the SgModel and load the checkpoint
-        model = SgModel(self.experiment_name, model_checkpoints_location='local')
-        model.build_model(new_torch_net, arch_params={'num_classes': 10},
-                          checkpoint_params={'strict_load': StrictLoad.NO_KEY_MATCHING,
-                                             'external_checkpoint_path': self.checkpoint_diff_keys_path,
-                                             'load_checkpoint': True})
+        # Build the Trainer and load the checkpoint
+        trainer = Trainer(self.experiment_name, model_checkpoints_location='local')
+        trainer.build_model(new_torch_net, arch_params={'num_classes': 10},
+                            checkpoint_params={'strict_load': StrictLoad.NO_KEY_MATCHING,
+                                               'external_checkpoint_path': self.checkpoint_diff_keys_path,
+                                               'load_checkpoint': True})
 
         # Assert the weights were loaded correctly
-        assert self.check_models_have_same_weights(model.net, self.original_torch_net)
+        assert self.check_models_have_same_weights(trainer.net, self.original_torch_net)
 
     def test_strict_load_no_key_matching_sg_checkpoint(self):
         # Define Model
@@ -147,15 +147,15 @@ class StrictLoadEnumTest(unittest.TestCase):
         # Make sure we initialized a model with different weights
         assert not self.check_models_have_same_weights(new_torch_net, self.original_torch_net)
 
-        # Build the SgModel and load the checkpoint
-        model = SgModel(self.experiment_name, model_checkpoints_location='local',
-                        ckpt_name='ckpt_latest_weights_only.pth')
-        model.build_model(new_torch_net, arch_params={'num_classes': 10},
-                          checkpoint_params={'strict_load': StrictLoad.NO_KEY_MATCHING,
-                                             'load_checkpoint': True})
+        # Build the Trainer and load the checkpoint
+        trainer = Trainer(self.experiment_name, model_checkpoints_location='local',
+                          ckpt_name='ckpt_latest_weights_only.pth')
+        trainer.build_model(new_torch_net, arch_params={'num_classes': 10},
+                            checkpoint_params={'strict_load': StrictLoad.NO_KEY_MATCHING,
+                                               'load_checkpoint': True})
 
         # Assert the weights were loaded correctly
-        assert self.check_models_have_same_weights(model.net, self.original_torch_net)
+        assert self.check_models_have_same_weights(trainer.net, self.original_torch_net)
 
 
 if __name__ == '__main__':

@@ -1,7 +1,7 @@
 import shutil
 import unittest
 import os
-from super_gradients import SgModel, \
+from super_gradients import Trainer, \
     ClassificationTestDatasetInterface, \
     SegmentationTestDatasetInterface, DetectionTestDatasetInterface
 from super_gradients.training.metrics import Accuracy, Top5
@@ -27,12 +27,12 @@ class TestWithoutTrainTest(unittest.TestCase):
 
     @staticmethod
     def get_classification_trainer(name=''):
-        model = SgModel(name, model_checkpoints_location='local')
+        trainer = Trainer(name, model_checkpoints_location='local')
         dataset_params = {"batch_size": 4}
         dataset = ClassificationTestDatasetInterface(dataset_params=dataset_params)
-        model.connect_dataset_interface(dataset)
-        model.build_model("resnet18_cifar")
-        return model
+        trainer.connect_dataset_interface(dataset)
+        trainer.build_model("resnet18_cifar")
+        return trainer
 
     @staticmethod
     def get_detection_trainer(name=''):
@@ -45,46 +45,46 @@ class TestWithoutTrainTest(unittest.TestCase):
                           "train_collate_fn": DetectionCollateFN(),
                           }
 
-        model = SgModel(name, model_checkpoints_location='local',
-                        multi_gpu=MultiGPUMode.OFF,
-                        post_prediction_callback=YoloPostPredictionCallback())
+        trainer = Trainer(name, model_checkpoints_location='local',
+                          multi_gpu=MultiGPUMode.OFF,
+                          post_prediction_callback=YoloPostPredictionCallback())
         dataset_interface = DetectionTestDatasetInterface(dataset_params=dataset_params)
-        model.connect_dataset_interface(dataset_interface, data_loader_num_workers=4)
-        model.build_model('yolox_s')
-        return model
+        trainer.connect_dataset_interface(dataset_interface, data_loader_num_workers=4)
+        trainer.build_model('yolox_s')
+        return trainer
 
     @staticmethod
     def get_segmentation_trainer(name=''):
         shelfnet_lw_arch_params = {"num_classes": 5, "load_checkpoint": False}
-        model = SgModel(name, model_checkpoints_location='local', multi_gpu=False)
+        trainer = Trainer(name, model_checkpoints_location='local', multi_gpu=False)
 
         dataset_interface = SegmentationTestDatasetInterface()
-        model.connect_dataset_interface(dataset_interface, data_loader_num_workers=8)
-        model.build_model('shelfnet34_lw', arch_params=shelfnet_lw_arch_params)
-        return model
+        trainer.connect_dataset_interface(dataset_interface, data_loader_num_workers=8)
+        trainer.build_model('shelfnet34_lw', arch_params=shelfnet_lw_arch_params)
+        return trainer
 
     def test_test_without_train(self):
-        model = self.get_classification_trainer(self.folder_names[0])
-        assert isinstance(model.test(silent_mode=True, test_metrics_list=[Accuracy(), Top5()]), tuple)
+        trainer = self.get_classification_trainer(self.folder_names[0])
+        assert isinstance(trainer.test(silent_mode=True, test_metrics_list=[Accuracy(), Top5()]), tuple)
 
-        model = self.get_detection_trainer(self.folder_names[1])
+        trainer = self.get_detection_trainer(self.folder_names[1])
 
-        test_metrics = [DetectionMetrics(post_prediction_callback=model.post_prediction_callback, num_cls=5)]
+        test_metrics = [DetectionMetrics(post_prediction_callback=trainer.post_prediction_callback, num_cls=5)]
 
-        assert isinstance(model.test(silent_mode=True, test_metrics_list=test_metrics), tuple)
+        assert isinstance(trainer.test(silent_mode=True, test_metrics_list=test_metrics), tuple)
 
-        model = self.get_segmentation_trainer(self.folder_names[2])
-        assert isinstance(model.test(silent_mode=True, test_metrics_list=[IoU(21), PixelAccuracy()]), tuple)
+        trainer = self.get_segmentation_trainer(self.folder_names[2])
+        assert isinstance(trainer.test(silent_mode=True, test_metrics_list=[IoU(21), PixelAccuracy()]), tuple)
 
     def test_test_on_valid_loader_without_train(self):
-        model = self.get_classification_trainer(self.folder_names[0])
-        assert isinstance(model.test(test_loader=model.valid_loader, silent_mode=True, test_metrics_list=[Accuracy(), Top5()]), tuple)
+        trainer = self.get_classification_trainer(self.folder_names[0])
+        assert isinstance(trainer.test(test_loader=trainer.valid_loader, silent_mode=True, test_metrics_list=[Accuracy(), Top5()]), tuple)
 
-        model = self.get_detection_trainer(self.folder_names[1])
+        trainer = self.get_detection_trainer(self.folder_names[1])
 
-        test_metrics = [DetectionMetrics(post_prediction_callback=model.post_prediction_callback, num_cls=5)]
+        test_metrics = [DetectionMetrics(post_prediction_callback=trainer.post_prediction_callback, num_cls=5)]
 
-        assert isinstance(model.test(test_loader=model.valid_loader, silent_mode=True, test_metrics_list=test_metrics), tuple)
+        assert isinstance(trainer.test(test_loader=trainer.valid_loader, silent_mode=True, test_metrics_list=test_metrics), tuple)
 
         model = self.get_segmentation_trainer(self.folder_names[2])
         assert isinstance(model.test(test_loader=model.valid_loader, silent_mode=True, test_metrics_list=[IoU(21), PixelAccuracy()]), tuple)

@@ -1,5 +1,5 @@
 import unittest
-from super_gradients import SgModel, \
+from super_gradients import Trainer, \
     ClassificationTestDatasetInterface
 from super_gradients.training.metrics import Accuracy, Top5
 from super_gradients.training.models import ResNet18
@@ -11,13 +11,13 @@ import shutil
 
 class SgTrainerLoggingTest(unittest.TestCase):
     def test_train_logging(self):
-        model = SgModel("test_train_with_full_log", model_checkpoints_location='local')
+        trainer = Trainer("test_train_with_full_log", model_checkpoints_location='local')
         dataset_params = {"batch_size": 10}
         dataset = ClassificationTestDatasetInterface(dataset_params=dataset_params)
-        model.connect_dataset_interface(dataset)
+        trainer.connect_dataset_interface(dataset)
 
         net = ResNet18(num_classes=5, arch_params={})
-        model.build_model(net)
+        trainer.build_model(net)
         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": "cross_entropy", "optimizer": "SGD",
                         "criterion_params": {}, "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9},
@@ -26,9 +26,9 @@ class SgTrainerLoggingTest(unittest.TestCase):
                         "greater_metric_to_watch_is_better": True,
                         "save_full_train_log": True}
 
-        model.train(train_params)
+        trainer.train(train_params)
 
-        logfile_path = model.log_file.replace('.txt', 'full_train_log.log')
+        logfile_path = trainer.log_file.replace('.txt', 'full_train_log.log')
         assert os.path.exists(logfile_path) and os.path.getsize(logfile_path) > 0
 
         root_logger_handlers = logging.root.handlers
@@ -41,7 +41,7 @@ class SgTrainerLoggingTest(unittest.TestCase):
         if os.path.exists(logs_dir_path):
             shutil.rmtree(logs_dir_path)
 
-        module_name = 'super_gradients.trainer.sg_model'
+        module_name = 'super_gradients.trainer.sg_trainer'
 
         _ = get_logger(module_name, training_log_path=None, logs_dir_path=logs_dir_path)
         root_logger_handlers = logging.root.handlers

@@ -1,5 +1,5 @@
 import unittest
-from super_gradients import SgModel, \
+from super_gradients import Trainer, \
     ClassificationTestDatasetInterface
 from super_gradients.training.metrics import Accuracy, Top5, ToyTestClassificationMetric
 from super_gradients.training.models import ResNet18
@@ -16,51 +16,52 @@ class TrainWithInitializedObjectsTest(unittest.TestCase):
     """
     Unit test for training with initialized objects passed as parameters.
     """
+
     def test_train_with_external_criterion(self):
-        model = SgModel("external_criterion_test", model_checkpoints_location='local')
+        trainer = Trainer("external_criterion_test", model_checkpoints_location='local')
         dataset_params = {"batch_size": 10}
         dataset = ClassificationTestDatasetInterface(dataset_params=dataset_params)
-        model.connect_dataset_interface(dataset)
+        trainer.connect_dataset_interface(dataset)
 
         net = ResNet18(num_classes=5, arch_params={})
-        model.build_model(net)
+        trainer.build_model(net)
         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}
-        model.train(train_params)
+        trainer.train(train_params)
 
     def test_train_with_external_optimizer(self):
-        model = SgModel("external_optimizer_test", model_checkpoints_location='local')
+        trainer = Trainer("external_optimizer_test", model_checkpoints_location='local')
         dataset_params = {"batch_size": 10}
         dataset = ClassificationTestDatasetInterface(dataset_params=dataset_params)
-        model.connect_dataset_interface(dataset)
+        trainer.connect_dataset_interface(dataset)
 
         net = ResNet18(num_classes=5, arch_params={})
         optimizer = SGD(params=net.parameters(), lr=0.1)
-        model.build_model(net)
+        trainer.build_model(net)
         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": "cross_entropy", "optimizer": optimizer,
                         "criterion_params": {}, "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9},
                         "train_metrics_list": [Accuracy(), Top5()], "valid_metrics_list": [Accuracy(), Top5()],
                         "loss_logging_items_names": ["Loss"], "metric_to_watch": "Accuracy",
                         "greater_metric_to_watch_is_better": True}
-        model.train(train_params)
+        trainer.train(train_params)
 
     def test_train_with_external_scheduler(self):
-        model = SgModel("external_scheduler_test", model_checkpoints_location='local')
+        trainer = Trainer("external_scheduler_test", model_checkpoints_location='local')
         dataset_params = {"batch_size": 10}
         dataset = ClassificationTestDatasetInterface(dataset_params=dataset_params)
-        model.connect_dataset_interface(dataset)
+        trainer.connect_dataset_interface(dataset)
 
         lr = 0.3
         net = ResNet18(num_classes=5, arch_params={})
         optimizer = SGD(params=net.parameters(), lr=lr)
         lr_scheduler = MultiStepLR(optimizer=optimizer, milestones=[1, 2], gamma=0.1)
         phase_callbacks = [LRSchedulerCallback(lr_scheduler, Phase.TRAIN_EPOCH_END)]
-        model.build_model(net)
+        trainer.build_model(net)
 
         train_params = {"max_epochs": 2, "phase_callbacks": phase_callbacks,
                         "lr_warmup_epochs": 0, "initial_lr": lr, "loss": "cross_entropy", "optimizer": optimizer,
@@ -68,18 +69,18 @@ class TrainWithInitializedObjectsTest(unittest.TestCase):
                         "train_metrics_list": [Accuracy(), Top5()], "valid_metrics_list": [Accuracy(), Top5()],
                         "loss_logging_items_names": ["Loss"], "metric_to_watch": "Accuracy",
                         "greater_metric_to_watch_is_better": True}
-        model.train(train_params)
+        trainer.train(train_params)
         assert lr_scheduler.get_last_lr()[0] == lr * 0.1 * 0.1
 
     def test_train_with_external_scheduler_class(self):
-        model = SgModel("external_scheduler_test", model_checkpoints_location='local')
+        trainer = Trainer("external_scheduler_test", model_checkpoints_location='local')
         dataset_params = {"batch_size": 10}
         dataset = ClassificationTestDatasetInterface(dataset_params=dataset_params)
-        model.connect_dataset_interface(dataset)
+        trainer.connect_dataset_interface(dataset)
 
         net = ResNet18(num_classes=5, arch_params={})
         optimizer = SGD  # a class - not an instance
-        model.build_model(net)
+        trainer.build_model(net)
 
         train_params = {"max_epochs": 2,
                         "lr_warmup_epochs": 0, "initial_lr": 0.3, "loss": "cross_entropy", "optimizer": optimizer,
@@ -87,20 +88,20 @@ class TrainWithInitializedObjectsTest(unittest.TestCase):
                         "train_metrics_list": [Accuracy(), Top5()], "valid_metrics_list": [Accuracy(), Top5()],
                         "loss_logging_items_names": ["Loss"], "metric_to_watch": "Accuracy",
                         "greater_metric_to_watch_is_better": True}
-        model.train(train_params)
+        trainer.train(train_params)
 
     def test_train_with_reduce_on_plateau(self):
-        model = SgModel("external_reduce_on_plateau_scheduler_test", model_checkpoints_location='local')
+        trainer = Trainer("external_reduce_on_plateau_scheduler_test", model_checkpoints_location='local')
         dataset_params = {"batch_size": 10}
         dataset = ClassificationTestDatasetInterface(dataset_params=dataset_params)
-        model.connect_dataset_interface(dataset)
+        trainer.connect_dataset_interface(dataset)
 
         lr = 0.3
         net = ResNet18(num_classes=5, arch_params={})
         optimizer = SGD(params=net.parameters(), lr=lr)
         lr_scheduler = ReduceLROnPlateau(optimizer=optimizer, patience=0)
         phase_callbacks = [LRSchedulerCallback(lr_scheduler, Phase.VALIDATION_EPOCH_END, "ToyTestClassificationMetric")]
-        model.build_model(net)
+        trainer.build_model(net)
 
         train_params = {"max_epochs": 2, "phase_callbacks": phase_callbacks,
                         "lr_warmup_epochs": 0, "initial_lr": lr, "loss": "cross_entropy", "optimizer": optimizer,
@@ -109,27 +110,27 @@ class TrainWithInitializedObjectsTest(unittest.TestCase):
                         "valid_metrics_list": [Accuracy(), Top5(), ToyTestClassificationMetric()],
                         "loss_logging_items_names": ["Loss"], "metric_to_watch": "Accuracy",
                         "greater_metric_to_watch_is_better": True}
-        model.train(train_params)
+        trainer.train(train_params)
         assert lr_scheduler._last_lr[0] == lr * 0.1
 
     def test_train_with_external_metric(self):
-        model = SgModel("external_metric_test", model_checkpoints_location='local')
+        trainer = Trainer("external_metric_test", model_checkpoints_location='local')
         dataset_params = {"batch_size": 10}
         dataset = ClassificationTestDatasetInterface(dataset_params=dataset_params)
-        model.connect_dataset_interface(dataset)
+        trainer.connect_dataset_interface(dataset)
 
         net = ResNet18(num_classes=5, arch_params={})
-        model.build_model(net)
+        trainer.build_model(net)
         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": "cross_entropy", "optimizer": "SGD",
                         "criterion_params": {}, "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9},
                         "train_metrics_list": [F1Score()], "valid_metrics_list": [F1Score()],
                         "loss_logging_items_names": ["Loss"], "metric_to_watch": "F1Score",
                         "greater_metric_to_watch_is_better": True}
-        model.train(train_params)
+        trainer.train(train_params)
 
     def test_train_with_external_dataloaders(self):
-        model = SgModel("external_data_loader_test", model_checkpoints_location='local')
+        trainer = Trainer("external_data_loader_test", model_checkpoints_location='local')
 
         batch_size = 5
         trainset = torch.utils.data.TensorDataset(torch.Tensor(np.random.random((10, 3, 32, 32))), torch.LongTensor(np.zeros((10))))
@@ -141,17 +142,17 @@ class TrainWithInitializedObjectsTest(unittest.TestCase):
         val_loader = torch.utils.data.DataLoader(valset, batch_size=batch_size)
 
         dataset_interface = DatasetInterface(train_loader=train_loader, val_loader=val_loader, classes=classes)
-        model.connect_dataset_interface(dataset_interface)
+        trainer.connect_dataset_interface(dataset_interface)
 
         net = ResNet18(num_classes=5, arch_params={})
-        model.build_model(net)
+        trainer.build_model(net)
         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": "cross_entropy", "optimizer": "SGD",
                         "criterion_params": {}, "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9},
                         "train_metrics_list": [F1Score()], "valid_metrics_list": [F1Score()],
                         "loss_logging_items_names": ["Loss"], "metric_to_watch": "F1Score",
                         "greater_metric_to_watch_is_better": True}
-        model.train(train_params)
+        trainer.train(train_params)
 
 
 if __name__ == '__main__':

@@ -1,5 +1,5 @@
 import unittest
-from super_gradients import SgModel, \
+from super_gradients import Trainer, \
     ClassificationTestDatasetInterface
 from super_gradients.training.metrics import Accuracy, Top5
 from super_gradients.training.models import ResNet18
@@ -11,13 +11,13 @@ class TrainWithPreciseBNTest(unittest.TestCase):
     """
 
     def test_train_with_precise_bn_explicit_size(self):
-        model = SgModel("test_train_with_precise_bn_explicit_size", model_checkpoints_location='local')
+        trainer = Trainer("test_train_with_precise_bn_explicit_size", model_checkpoints_location='local')
         dataset_params = {"batch_size": 10}
         dataset = ClassificationTestDatasetInterface(dataset_params=dataset_params)
-        model.connect_dataset_interface(dataset)
+        trainer.connect_dataset_interface(dataset)
 
         net = ResNet18(num_classes=5, arch_params={})
-        model.build_model(net)
+        trainer.build_model(net)
         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": "cross_entropy", "optimizer": "SGD",
                         "criterion_params": {}, "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9},
@@ -25,16 +25,16 @@ class TrainWithPreciseBNTest(unittest.TestCase):
                         "loss_logging_items_names": ["Loss"], "metric_to_watch": "Accuracy",
                         "greater_metric_to_watch_is_better": True,
                         "precise_bn": True, "precise_bn_batch_size": 100}
-        model.train(train_params)
+        trainer.train(train_params)
 
     def test_train_with_precise_bn_implicit_size(self):
-        model = SgModel("test_train_with_precise_bn_implicit_size", model_checkpoints_location='local')
+        trainer = Trainer("test_train_with_precise_bn_implicit_size", model_checkpoints_location='local')
         dataset_params = {"batch_size": 10}
         dataset = ClassificationTestDatasetInterface(dataset_params=dataset_params)
-        model.connect_dataset_interface(dataset)
+        trainer.connect_dataset_interface(dataset)
 
         net = ResNet18(num_classes=5, arch_params={})
-        model.build_model(net)
+        trainer.build_model(net)
         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": "cross_entropy", "optimizer": "SGD",
                         "criterion_params": {}, "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9},
@@ -42,7 +42,7 @@ class TrainWithPreciseBNTest(unittest.TestCase):
                         "loss_logging_items_names": ["Loss"], "metric_to_watch": "Accuracy",
                         "greater_metric_to_watch_is_better": True,
                         "precise_bn": True}
-        model.train(train_params)
+        trainer.train(train_params)
 
 
 if __name__ == '__main__':

@@ -1,5 +1,5 @@
 import unittest
-from super_gradients.training import SgModel
+from super_gradients.training import Trainer
 from super_gradients.training.metrics import Accuracy
 from super_gradients.training.datasets import ClassificationTestDatasetInterface
 from super_gradients.training.models import LeNet
@@ -33,9 +33,9 @@ class UpdateParamGroupsTest(unittest.TestCase):
     def test_lr_scheduling_with_update_param_groups(self):
         # Define Model
         net = TestNet()
-        model = SgModel("lr_warmup_test", model_checkpoints_location='local')
-        model.connect_dataset_interface(self.dataset)
-        model.build_model(net, arch_params=self.arch_params)
+        trainer = Trainer("lr_warmup_test", model_checkpoints_location='local')
+        trainer.connect_dataset_interface(self.dataset)
+        trainer.build_model(net, arch_params=self.arch_params)
 
         lrs = []
         phase_callbacks = [TestLRCallback(lr_placeholder=lrs)]
@@ -53,6 +53,6 @@ class UpdateParamGroupsTest(unittest.TestCase):
                         }
 
         expected_lrs = np.array([0.1, 0.2, 0.3])
-        model.train(train_params)
+        trainer.train(train_params)
 
         self.assertTrue(np.allclose(np.array(lrs), expected_lrs, rtol=0.0000001))

@@ -1,7 +1,7 @@
 import unittest
 from super_gradients.training.utils.utils import HpmStruct
 from super_gradients.training.datasets.dataset_interfaces.dataset_interface import ClassificationTestDatasetInterface
-from super_gradients import SgModel
+from super_gradients import Trainer
 from super_gradients.training.metrics import Accuracy, Top5
 
 
@@ -21,10 +21,10 @@ class TestViT(unittest.TestCase):
         """
         Validate vit_base
         """
-        model = SgModel("test_vit_base", device='cpu')
-        model.connect_dataset_interface(self.dataset, data_loader_num_workers=8)
-        model.build_model('vit_base', load_checkpoint=False)
-        model.train(training_params=self.train_params)
+        trainer = Trainer("test_vit_base", device='cpu')
+        trainer.connect_dataset_interface(self.dataset, data_loader_num_workers=8)
+        trainer.build_model('vit_base', load_checkpoint=False)
+        trainer.train(training_params=self.train_params)
 
 
 if __name__ == '__main__':