@@ -5,14 +5,10 @@ import re
 from super_gradients import (
     SgModel,
     ClassificationTestDatasetInterface,
-    DetectionTestDatasetInterface,
     SegmentationTestDatasetInterface,
 )
 from super_gradients.training.utils.callbacks import ModelConversionCheckCallback
-from super_gradients.training.utils.detection_utils import Anchors
-from super_gradients.training.models.detection_models.yolo_base import YoloPostPredictionCallback
 from super_gradients.training.metrics import Accuracy, Top5, IoU
-from super_gradients.training.metrics.detection_metrics import DetectionMetrics
 from super_gradients.training.losses.stdc_loss import STDCLoss
 from super_gradients.training.losses.ddrnet_loss import DDRNetLoss
 
@@ -48,6 +44,7 @@ CLASSIFICATION = ["efficientnet_b0", "regnetY200", "regnetY400", "regnetY600", "
 SEMANTIC_SEGMENTATION = ["ddrnet_23", "stdc1_seg", "stdc2_seg", "regseg48"]
 # TODO: ADD YOLOX ARCHITECTURES AND TESTS
 
+
 class ConversionCallbackTest(unittest.TestCase):
     def test_classification_architectures(self):
         for architecture in CLASSIFICATION:

@@ -80,8 +80,7 @@ class PretrainedModelsTest(unittest.TestCase):
                                                      "metric_to_watch": "Accuracy",
                                                      "greater_metric_to_watch_is_better": True}
         self.coco_pretrained_ckpt_params = {"pretrained_weights": "coco"}
-        self.coco_pretrained_arch_params = {
-                                            'ssd_lite_mobilenet_v2': {'num_classes': 80},
+        self.coco_pretrained_arch_params = {'ssd_lite_mobilenet_v2': {'num_classes': 80},
                                             'coco_ssd_mobilenet_v1': {'num_classes': 80}}
         self.coco_pretrained_ckpt_params = {"pretrained_weights": "coco"}
         self.coco_dataset = {
@@ -280,7 +279,8 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = SgModel('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"], checkpoint_params=self.imagenet_pretrained_ckpt_params)
+        trainer.build_model("resnet50", arch_params=self.imagenet_pretrained_arch_params["resnet"],
+                            checkpoint_params=self.imagenet_pretrained_ckpt_params)
         res = trainer.test(test_loader=self.imagenet_dataset.val_loader, test_metrics_list=[Accuracy()],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["resnet50"], delta=0.001)
@@ -289,14 +289,16 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = SgModel('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"], checkpoint_params=self.imagenet_pretrained_ckpt_params)
+        trainer.build_model("resnet50", arch_params=self.imagenet_pretrained_arch_params["resnet"],
+                            checkpoint_params=self.imagenet_pretrained_ckpt_params)
         trainer.train(training_params=self.transfer_classification_train_params)
 
     def test_pretrained_resnet34_imagenet(self):
         trainer = SgModel('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"], checkpoint_params=self.imagenet_pretrained_ckpt_params)
+        trainer.build_model("resnet34", arch_params=self.imagenet_pretrained_arch_params["resnet"],
+                            checkpoint_params=self.imagenet_pretrained_ckpt_params)
         res = trainer.test(test_loader=self.imagenet_dataset.val_loader, test_metrics_list=[Accuracy()],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["resnet34"], delta=0.001)
@@ -305,14 +307,16 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = SgModel('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"], checkpoint_params=self.imagenet_pretrained_ckpt_params)
+        trainer.build_model("resnet34", arch_params=self.imagenet_pretrained_arch_params["resnet"],
+                            checkpoint_params=self.imagenet_pretrained_ckpt_params)
         trainer.train(training_params=self.transfer_classification_train_params)
 
     def test_pretrained_resnet18_imagenet(self):
         trainer = SgModel('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"], checkpoint_params=self.imagenet_pretrained_ckpt_params)
+        trainer.build_model("resnet18", arch_params=self.imagenet_pretrained_arch_params["resnet"],
+                            checkpoint_params=self.imagenet_pretrained_ckpt_params)
         res = trainer.test(test_loader=self.imagenet_dataset.val_loader, test_metrics_list=[Accuracy()],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["resnet18"], delta=0.001)
@@ -321,14 +325,16 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = SgModel('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"], checkpoint_params=self.imagenet_pretrained_ckpt_params)
+        trainer.build_model("resnet18", arch_params=self.imagenet_pretrained_arch_params["resnet"],
+                            checkpoint_params=self.imagenet_pretrained_ckpt_params)
         trainer.train(training_params=self.transfer_classification_train_params)
 
     def test_pretrained_regnetY800_imagenet(self):
         trainer = SgModel('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"], checkpoint_params=self.imagenet_pretrained_ckpt_params)
+        trainer.build_model("regnetY800", arch_params=self.imagenet_pretrained_arch_params["regnet"],
+                            checkpoint_params=self.imagenet_pretrained_ckpt_params)
         res = trainer.test(test_loader=self.imagenet_dataset.val_loader, test_metrics_list=[Accuracy()],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["regnetY800"], delta=0.001)
@@ -337,14 +343,16 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = SgModel('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"], checkpoint_params=self.imagenet_pretrained_ckpt_params)
+        trainer.build_model("regnetY800", arch_params=self.imagenet_pretrained_arch_params["regnet"],
+                            checkpoint_params=self.imagenet_pretrained_ckpt_params)
         trainer.train(training_params=self.transfer_classification_train_params)
 
     def test_pretrained_regnetY600_imagenet(self):
         trainer = SgModel('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"], checkpoint_params=self.imagenet_pretrained_ckpt_params)
+        trainer.build_model("regnetY600", arch_params=self.imagenet_pretrained_arch_params["regnet"],
+                            checkpoint_params=self.imagenet_pretrained_ckpt_params)
         res = trainer.test(test_loader=self.imagenet_dataset.val_loader, test_metrics_list=[Accuracy()],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["regnetY600"], delta=0.001)
@@ -353,14 +361,16 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = SgModel('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"], checkpoint_params=self.imagenet_pretrained_ckpt_params)
+        trainer.build_model("regnetY600", arch_params=self.imagenet_pretrained_arch_params["regnet"],
+                            checkpoint_params=self.imagenet_pretrained_ckpt_params)
         trainer.train(training_params=self.transfer_classification_train_params)
 
     def test_pretrained_regnetY400_imagenet(self):
         trainer = SgModel('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"], checkpoint_params=self.imagenet_pretrained_ckpt_params)
+        trainer.build_model("regnetY400", arch_params=self.imagenet_pretrained_arch_params["regnet"],
+                            checkpoint_params=self.imagenet_pretrained_ckpt_params)
         res = trainer.test(test_loader=self.imagenet_dataset.val_loader, test_metrics_list=[Accuracy()],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["regnetY400"], delta=0.001)
@@ -369,14 +379,16 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = SgModel('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"], checkpoint_params=self.imagenet_pretrained_ckpt_params)
+        trainer.build_model("regnetY400", arch_params=self.imagenet_pretrained_arch_params["regnet"],
+                            checkpoint_params=self.imagenet_pretrained_ckpt_params)
         trainer.train(training_params=self.transfer_classification_train_params)
 
     def test_pretrained_regnetY200_imagenet(self):
         trainer = SgModel('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"], checkpoint_params=self.imagenet_pretrained_ckpt_params)
+        trainer.build_model("regnetY200", arch_params=self.imagenet_pretrained_arch_params["regnet"],
+                            checkpoint_params=self.imagenet_pretrained_ckpt_params)
         res = trainer.test(test_loader=self.imagenet_dataset.val_loader, test_metrics_list=[Accuracy()],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["regnetY200"], delta=0.001)
@@ -385,14 +397,16 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = SgModel('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"], checkpoint_params=self.imagenet_pretrained_ckpt_params)
+        trainer.build_model("regnetY200", arch_params=self.imagenet_pretrained_arch_params["regnet"],
+                            checkpoint_params=self.imagenet_pretrained_ckpt_params)
         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',
                           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"], checkpoint_params=self.imagenet_pretrained_ckpt_params)
+        trainer.build_model("repvgg_a0", arch_params=self.imagenet_pretrained_arch_params["repvgg_a0"],
+                            checkpoint_params=self.imagenet_pretrained_ckpt_params)
         res = trainer.test(test_loader=self.imagenet_dataset.val_loader, test_metrics_list=[Accuracy()],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["repvgg_a0"], delta=0.001)
@@ -401,14 +415,16 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = SgModel('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"], checkpoint_params=self.imagenet_pretrained_ckpt_params)
+        trainer.build_model("repvgg_a0", arch_params=self.imagenet_pretrained_arch_params["repvgg_a0"],
+                            checkpoint_params=self.imagenet_pretrained_ckpt_params)
         trainer.train(training_params=self.transfer_classification_train_params)
 
     def test_pretrained_regseg48_cityscapes(self):
         trainer = SgModel('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"], checkpoint_params=self.cityscapes_pretrained_ckpt_params)
+        trainer.build_model("regseg48", arch_params=self.cityscapes_pretrained_arch_params["regseg48"],
+                            checkpoint_params=self.cityscapes_pretrained_ckpt_params)
         res = trainer.test(test_loader=self.cityscapes_dataset.val_loader,
                            test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
                            metrics_progress_verbose=True)[0].cpu().item()
@@ -418,14 +434,16 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = SgModel('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"], checkpoint_params=self.cityscapes_pretrained_ckpt_params)
+        trainer.build_model("regseg48", arch_params=self.cityscapes_pretrained_arch_params["regseg48"],
+                            checkpoint_params=self.cityscapes_pretrained_ckpt_params)
         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',
                           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"], checkpoint_params=self.cityscapes_pretrained_ckpt_params)
+        trainer.build_model("ddrnet_23", arch_params=self.cityscapes_pretrained_arch_params["ddrnet_23"],
+                            checkpoint_params=self.cityscapes_pretrained_ckpt_params)
         res = trainer.test(test_loader=self.cityscapes_dataset.val_loader,
                            test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
                            metrics_progress_verbose=True)[0].cpu().item()
@@ -435,7 +453,8 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = SgModel('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"], checkpoint_params=self.cityscapes_pretrained_ckpt_params)
+        trainer.build_model("ddrnet_23_slim", arch_params=self.cityscapes_pretrained_arch_params["ddrnet_23"],
+                            checkpoint_params=self.cityscapes_pretrained_ckpt_params)
         res = trainer.test(test_loader=self.cityscapes_dataset.val_loader,
                            test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
                            metrics_progress_verbose=True)[0].cpu().item()
@@ -445,14 +464,16 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = SgModel('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"], checkpoint_params=self.cityscapes_pretrained_ckpt_params)
+        trainer.build_model("ddrnet_23", arch_params=self.cityscapes_pretrained_arch_params["ddrnet_23"],
+                            checkpoint_params=self.cityscapes_pretrained_ckpt_params)
         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',
                           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"], checkpoint_params=self.cityscapes_pretrained_ckpt_params)
+        trainer.build_model("ddrnet_23_slim", arch_params=self.cityscapes_pretrained_arch_params["ddrnet_23"],
+                            checkpoint_params=self.cityscapes_pretrained_ckpt_params)
         trainer.train(training_params=self.ddrnet_transfer_segmentation_train_params)
 
     def test_pretrained_coco_segmentation_subclass_pretrained_shelfnet34_lw(self):
@@ -460,7 +481,8 @@ class PretrainedModelsTest(unittest.TestCase):
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.coco_segmentation_dataset, data_loader_num_workers=8)
         trainer.build_model("shelfnet34_lw",
-                            arch_params=self.coco_segmentation_subclass_pretrained_arch_params["shelfnet34_lw"], checkpoint_params=self.coco_segmentation_subclass_pretrained_ckpt_params)
+                            arch_params=self.coco_segmentation_subclass_pretrained_arch_params["shelfnet34_lw"],
+                            checkpoint_params=self.coco_segmentation_subclass_pretrained_ckpt_params)
         res = trainer.test(test_loader=self.coco_segmentation_dataset.val_loader, test_metrics_list=[IoU(21)],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.coco_segmentation_subclass_pretrained_mious["shelfnet34_lw"], delta=0.001)
@@ -469,7 +491,8 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = SgModel('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"], checkpoint_params=self.imagenet_pretrained_ckpt_params)
+        trainer.build_model("efficientnet_b0", arch_params=self.imagenet_pretrained_arch_params["efficientnet_b0"],
+                            checkpoint_params=self.imagenet_pretrained_ckpt_params)
         res = trainer.test(test_loader=self.imagenet_dataset.val_loader, test_metrics_list=[Accuracy()],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["efficientnet_b0"], delta=0.001)
@@ -478,7 +501,8 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = SgModel('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"], checkpoint_params=self.imagenet_pretrained_ckpt_params)
+        trainer.build_model("efficientnet_b0", arch_params=self.imagenet_pretrained_arch_params["efficientnet_b0"],
+                            checkpoint_params=self.imagenet_pretrained_ckpt_params)
         trainer.train(training_params=self.transfer_classification_train_params)
 
     def test_pretrained_ssd_lite_mobilenet_v2_coco(self):
@@ -517,7 +541,7 @@ class PretrainedModelsTest(unittest.TestCase):
         ssd_post_prediction_callback = SSDPostPredictCallback()
         res = trainer.test(test_loader=self.coco_dataset['ssd_mobilenet'].val_loader,
                            test_metrics_list=[DetectionMetrics(post_prediction_callback=ssd_post_prediction_callback,
-                                                                 num_cls=len(self.coco_dataset['ssd_mobilenet'].coco_classes))],
+                                                               num_cls=len(self.coco_dataset['ssd_mobilenet'].coco_classes))],
                            metrics_progress_verbose=True)[2]
         self.assertAlmostEqual(res, self.coco_pretrained_maps["coco_ssd_mobilenet_v1"], delta=0.001)
 
@@ -529,8 +553,8 @@ class PretrainedModelsTest(unittest.TestCase):
                             checkpoint_params=self.coco_pretrained_ckpt_params)
         res = trainer.test(test_loader=self.coco_dataset['yolox'].val_loader,
                            test_metrics_list=[DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(),
-                                                                 num_cls=80,
-                                                                 normalize_targets=True)])[2]
+                                                               num_cls=80,
+                                                               normalize_targets=True)])[2]
         self.assertAlmostEqual(res, self.coco_pretrained_maps["yolox_s"], delta=0.001)
 
     def test_pretrained_yolox_m_coco(self):
@@ -541,8 +565,8 @@ class PretrainedModelsTest(unittest.TestCase):
                             checkpoint_params=self.coco_pretrained_ckpt_params)
         res = trainer.test(test_loader=self.coco_dataset['yolox'].val_loader,
                            test_metrics_list=[DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(),
-                                                                 num_cls=80,
-                                                                 normalize_targets=True)])[2]
+                                                               num_cls=80,
+                                                               normalize_targets=True)])[2]
         self.assertAlmostEqual(res, self.coco_pretrained_maps["yolox_m"], delta=0.001)
 
     def test_pretrained_yolox_l_coco(self):
@@ -553,8 +577,8 @@ class PretrainedModelsTest(unittest.TestCase):
                             checkpoint_params=self.coco_pretrained_ckpt_params)
         res = trainer.test(test_loader=self.coco_dataset['yolox'].val_loader,
                            test_metrics_list=[DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(),
-                                                                 num_cls=80,
-                                                                 normalize_targets=True)])[2]
+                                                               num_cls=80,
+                                                               normalize_targets=True)])[2]
         self.assertAlmostEqual(res, self.coco_pretrained_maps["yolox_l"], delta=0.001)
 
     def test_pretrained_yolox_n_coco(self):
@@ -565,8 +589,8 @@ class PretrainedModelsTest(unittest.TestCase):
                             checkpoint_params=self.coco_pretrained_ckpt_params)
         res = trainer.test(test_loader=self.coco_dataset['yolox'].val_loader,
                            test_metrics_list=[DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(),
-                                                                 num_cls=80,
-                                                                 normalize_targets=True)])[2]
+                                                               num_cls=80,
+                                                               normalize_targets=True)])[2]
         self.assertAlmostEqual(res, self.coco_pretrained_maps["yolox_n"], delta=0.001)
 
     def test_pretrained_yolox_t_coco(self):
@@ -577,8 +601,8 @@ class PretrainedModelsTest(unittest.TestCase):
                             checkpoint_params=self.coco_pretrained_ckpt_params)
         res = trainer.test(test_loader=self.coco_dataset['yolox'].val_loader,
                            test_metrics_list=[DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(),
-                                                                 num_cls=80,
-                                                                 normalize_targets=True)])[2]
+                                                               num_cls=80,
+                                                               normalize_targets=True)])[2]
         self.assertAlmostEqual(res, self.coco_pretrained_maps["yolox_t"], delta=0.001)
 
     def test_transfer_learning_mobilenet_v3_large_imagenet(self):
@@ -586,14 +610,16 @@ class PretrainedModelsTest(unittest.TestCase):
                           model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
-        trainer.build_model("mobilenet_v3_large", arch_params=self.imagenet_pretrained_arch_params["mobilenet"], checkpoint_params=self.imagenet_pretrained_ckpt_params)
+        trainer.build_model("mobilenet_v3_large", arch_params=self.imagenet_pretrained_arch_params["mobilenet"],
+                            checkpoint_params=self.imagenet_pretrained_ckpt_params)
         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',
                           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"], checkpoint_params=self.imagenet_pretrained_ckpt_params)
+        trainer.build_model("mobilenet_v3_large", arch_params=self.imagenet_pretrained_arch_params["mobilenet"],
+                            checkpoint_params=self.imagenet_pretrained_ckpt_params)
         res = trainer.test(test_loader=self.imagenet_dataset.val_loader, test_metrics_list=[Accuracy()],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["mobilenet_v3_large"], delta=0.001)
@@ -603,14 +629,16 @@ class PretrainedModelsTest(unittest.TestCase):
                           model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
-        trainer.build_model("mobilenet_v3_small", arch_params=self.imagenet_pretrained_arch_params["mobilenet"], checkpoint_params=self.imagenet_pretrained_ckpt_params)
+        trainer.build_model("mobilenet_v3_small", arch_params=self.imagenet_pretrained_arch_params["mobilenet"],
+                            checkpoint_params=self.imagenet_pretrained_ckpt_params)
         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',
                           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"], checkpoint_params=self.imagenet_pretrained_ckpt_params)
+        trainer.build_model("mobilenet_v3_small", arch_params=self.imagenet_pretrained_arch_params["mobilenet"],
+                            checkpoint_params=self.imagenet_pretrained_ckpt_params)
         res = trainer.test(test_loader=self.imagenet_dataset.val_loader, test_metrics_list=[Accuracy()],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["mobilenet_v3_small"], delta=0.001)
@@ -620,14 +648,16 @@ class PretrainedModelsTest(unittest.TestCase):
                           model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)
         trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
-        trainer.build_model("mobilenet_v2", arch_params=self.imagenet_pretrained_arch_params["mobilenet"], checkpoint_params=self.imagenet_pretrained_ckpt_params)
+        trainer.build_model("mobilenet_v2", arch_params=self.imagenet_pretrained_arch_params["mobilenet"],
+                            checkpoint_params=self.imagenet_pretrained_ckpt_params)
         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',
                           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"], checkpoint_params=self.imagenet_pretrained_ckpt_params)
+        trainer.build_model("mobilenet_v2", arch_params=self.imagenet_pretrained_arch_params["mobilenet"],
+                            checkpoint_params=self.imagenet_pretrained_ckpt_params)
         res = trainer.test(test_loader=self.imagenet_dataset.val_loader, test_metrics_list=[Accuracy()],
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["mobilenet_v2"], delta=0.001)
@@ -636,7 +666,8 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = SgModel('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"], checkpoint_params=self.cityscapes_pretrained_ckpt_params)
+        trainer.build_model("stdc1_seg50", arch_params=self.cityscapes_pretrained_arch_params["stdc"],
+                            checkpoint_params=self.cityscapes_pretrained_ckpt_params)
         res = trainer.test(test_loader=self.cityscapes_dataset_rescaled50.val_loader,
                            test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
                            metrics_progress_verbose=True)[0].cpu().item()
@@ -646,14 +677,16 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = SgModel('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"], checkpoint_params=self.cityscapes_pretrained_ckpt_params)
+        trainer.build_model("stdc1_seg50", arch_params=self.cityscapes_pretrained_arch_params["stdc"],
+                            checkpoint_params=self.cityscapes_pretrained_ckpt_params)
         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',
                           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"], checkpoint_params=self.cityscapes_pretrained_ckpt_params)
+        trainer.build_model("stdc1_seg75", arch_params=self.cityscapes_pretrained_arch_params["stdc"],
+                            checkpoint_params=self.cityscapes_pretrained_ckpt_params)
         res = trainer.test(test_loader=self.cityscapes_dataset_rescaled75.val_loader,
                            test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
                            metrics_progress_verbose=True)[0].cpu().item()
@@ -663,14 +696,16 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = SgModel('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"], checkpoint_params=self.cityscapes_pretrained_ckpt_params)
+        trainer.build_model("stdc1_seg75", arch_params=self.cityscapes_pretrained_arch_params["stdc"],
+                            checkpoint_params=self.cityscapes_pretrained_ckpt_params)
         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',
                           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"], checkpoint_params=self.cityscapes_pretrained_ckpt_params)
+        trainer.build_model("stdc2_seg50", arch_params=self.cityscapes_pretrained_arch_params["stdc"],
+                            checkpoint_params=self.cityscapes_pretrained_ckpt_params)
         res = trainer.test(test_loader=self.cityscapes_dataset_rescaled50.val_loader,
                            test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
                            metrics_progress_verbose=True)[0].cpu().item()
@@ -680,14 +715,16 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = SgModel('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"], checkpoint_params=self.cityscapes_pretrained_ckpt_params)
+        trainer.build_model("stdc2_seg50", arch_params=self.cityscapes_pretrained_arch_params["stdc"],
+                            checkpoint_params=self.cityscapes_pretrained_ckpt_params)
         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',
                           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"], checkpoint_params=self.cityscapes_pretrained_ckpt_params)
+        trainer.build_model("stdc2_seg75", arch_params=self.cityscapes_pretrained_arch_params["stdc"],
+                            checkpoint_params=self.cityscapes_pretrained_ckpt_params)
         res = trainer.test(test_loader=self.cityscapes_dataset_rescaled75.val_loader,
                            test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
                            metrics_progress_verbose=True)[0].cpu().item()
@@ -697,7 +734,8 @@ class PretrainedModelsTest(unittest.TestCase):
         trainer = SgModel('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"], checkpoint_params=self.cityscapes_pretrained_ckpt_params)
+        trainer.build_model("stdc2_seg75", arch_params=self.cityscapes_pretrained_arch_params["stdc"],
+                            checkpoint_params=self.cityscapes_pretrained_ckpt_params)
         trainer.train(training_params=self.stdc_transfer_segmentation_train_params)
 
     def test_transfer_learning_vit_base_imagenet21k(self):

@@ -40,7 +40,7 @@ class TestDatasetStatisticsTensorboardLogger(unittest.TestCase):
                            "lr_mode": "cosine",
                            "initial_lr": 0.01,
                            "loss": "yolox_loss",
-                           "criterion_params":{"strides": [8, 16, 32], "num_cls": 80},
+                           "criterion_params": {"strides": [8, 16, 32], "num_cls": 80},
                            "dataset_statistics": True,
                            "launch_tensorboard": True,
                            "valid_metrics_list": [DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(),

@@ -75,8 +75,6 @@ class TestUtil(unittest.TestCase):
         self.assertIsInstance(model._modules['0'], nn.Conv2d)
         self.assertIsInstance(model._modules['1'], nn.Conv2d)
 
-
-
     def test_fuse_conv_bn_on_toy_models(self):
 
         class Toy(nn.Module):

@@ -124,7 +124,7 @@ class LRWarmupTest(unittest.TestCase):
                         "train_metrics_list": [Accuracy()], "valid_metrics_list": [Accuracy()],
                         "loss_logging_items_names": ["Loss"], "metric_to_watch": "Accuracy",
                         "greater_metric_to_watch_is_better": True, "ema": False, "phase_callbacks": phase_callbacks,
-                        "warmup_mode": ExponentialWarmupLRCallback, "initial_lr": 1.,  "warmup_initial_lr": 0.1}
+                        "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)

@@ -149,7 +149,7 @@ class MaskAttentionLossTest(unittest.TestCase):
         """
         predict = torch.randn(self.batch, self.num_classes, self.img_size, self.img_size)
         target = torch.randint(0, self.num_classes,
-                               size=(self.batch, self.num_classes,  self.img_size, self.img_size)).float()
+                               size=(self.batch, self.num_classes, self.img_size, self.img_size)).float()
 
         loss_weigths = [1., 0.5]
         ce_crit = nn.BCEWithLogitsLoss(reduction="none")

@@ -49,11 +49,13 @@ class PhaseContextTest(unittest.TestCase):
 
                 if phase_callback.phase == Phase.VALIDATION_BATCH_END:
                     self.assertTrue(phase_callback.context.epoch == 2)
-                    self.assertTrue(isinstance(phase_callback.context.metrics_compute_fn, MetricCollection) and hasattr(phase_callback.context.metrics_compute_fn, "Top5"))
+                    self.assertTrue(isinstance(phase_callback.context.metrics_compute_fn, MetricCollection) and
+                                    hasattr(phase_callback.context.metrics_compute_fn, "Top5"))
 
                 else:
                     self.assertTrue(phase_callback.context.epoch == 1)
-                    self.assertTrue(isinstance(phase_callback.context.metrics_compute_fn, MetricCollection) and hasattr(phase_callback.context.metrics_compute_fn, "Accuracy"))
+                    self.assertTrue(isinstance(phase_callback.context.metrics_compute_fn, MetricCollection) and
+                                    hasattr(phase_callback.context.metrics_compute_fn, "Accuracy"))
 
         if phase_callback.phase in [Phase.TRAIN_EPOCH_END, Phase.VALIDATION_EPOCH_END]:
             self.assertTrue(phase_callback.context.batch_idx is None)
@@ -67,7 +69,9 @@ class PhaseContextTest(unittest.TestCase):
             self.assertTrue(phase_callback.context.epoch == 1)
 
             # EPOCH END PHASES USE THE SAME CONTEXT, WHICH IS UPDATED- SO VALID METRICS DICT SHOULD BE PRESENT
-            self.assertTrue(isinstance(phase_callback.context.metrics_dict, dict) and "Loss" in phase_callback.context.metrics_dict.keys() and "Top5" in phase_callback.context.metrics_dict.keys())
+            self.assertTrue(isinstance(phase_callback.context.metrics_dict, dict))
+            self.assertTrue("Loss" in phase_callback.context.metrics_dict.keys())
+            self.assertTrue("Top5" in phase_callback.context.metrics_dict.keys())
 
 
 if __name__ == '__main__':

@@ -17,10 +17,8 @@ class TestSgLibraryDatasets(TestCase):
 
     def test_get_dateset_with_invalid_dataset_name_raises_exception(self):
         with self.assertRaises(ValueError):
-            self.sg_library_datasets.get_dataset('classification',
-                                                   'cifar_1000000')
+            self.sg_library_datasets.get_dataset('classification', 'cifar_1000000')
 
     def test_get_dateset_with_invalid_dl_task_raises_exception(self):
         with self.assertRaises(ValueError):
-            self.sg_library_datasets.get_dataset('classification_of_something_that_deci_does_not_support',
-                                                   'cifar_100')
+            self.sg_library_datasets.get_dataset('classification_of_something_that_deci_does_not_support', 'cifar_100')

@@ -7,9 +7,7 @@ from super_gradients import SgModel, \
 from super_gradients.training.metrics import Accuracy, Top5
 from super_gradients.training import MultiGPUMode
 from super_gradients.training.models.detection_models.yolo_base import YoloPostPredictionCallback
-from super_gradients.training.utils.detection_utils import base_detection_collate_fn, DetectionCollateFN
-from super_gradients.training.datasets.datasets_utils import ComposedCollateFunction, MultiScaleCollateFunction
-from super_gradients.training.utils.detection_utils import YoloV3NonMaxSuppression
+from super_gradients.training.utils.detection_utils import DetectionCollateFN
 from super_gradients.training.metrics.detection_metrics import DetectionMetrics
 from super_gradients.training.metrics.segmentation_metrics import PixelAccuracy, IoU
 

@@ -78,7 +78,7 @@ class TrainWithInitializedObjectsTest(unittest.TestCase):
         model.connect_dataset_interface(dataset)
 
         net = ResNet18(num_classes=5, arch_params={})
-        optimizer = SGD # a class - not an instance
+        optimizer = SGD  # a class - not an instance
         model.build_model(net)
 
         train_params = {"max_epochs": 2,