@@ -11,7 +11,6 @@ import os
 import shutil
 from super_gradients.training.utils.ssd_utils import SSDPostPredictCallback
 from super_gradients.training.models.detection_models.ssd import DEFAULT_SSD_LITE_MOBILENET_V2_ARCH_PARAMS
-import torchvision.transforms as transforms
 from super_gradients.training.losses.ddrnet_loss import DDRNetLoss
 from super_gradients.training.metrics import DetectionMetrics
 from super_gradients.training.transforms.transforms import Rescale
@@ -232,7 +231,9 @@ class PretrainedModelsTest(unittest.TestCase):
         self.cityscapes_pretrained_arch_params = {
             "ddrnet_23": {"aux_head": True, "sync_bn": True},
             "regseg48": {},
-            "stdc": {"use_aux_heads": True, "aux_head": True}}
+            "stdc": {"use_aux_heads": True, "aux_head": True},
+            "pplite_seg": {"use_aux_heads": True},
+        }
 
         self.cityscapes_pretrained_ckpt_params = {"pretrained_weights": "cityscapes"}
         self.cityscapes_pretrained_mious = {"ddrnet_23": 0.8026,
@@ -241,7 +242,11 @@ class PretrainedModelsTest(unittest.TestCase):
                                             "stdc1_seg75": 0.7687,
                                             "stdc2_seg50": 0.7644,
                                             "stdc2_seg75": 0.7893,
-                                            "regseg48": 0.7815}
+                                            "regseg48": 0.7815,
+                                            "pp_lite_t_seg50": 0.7492,
+                                            "pp_lite_t_seg75": 0.7756,
+                                            "pp_lite_b_seg50": 0.7648,
+                                            "pp_lite_b_seg75": 0.7852}
 
         self.cityscapes_dataset = CityscapesDatasetInterface(dataset_params={
             "batch_size": 3,
@@ -249,22 +254,22 @@ class PretrainedModelsTest(unittest.TestCase):
             "dataset_dir": "/data/cityscapes/",
             "crop_size": 1024,
             "img_size": 1024,
-            "image_mask_transforms_aug": transforms.Compose([]),
-            "image_mask_transforms": transforms.Compose([])  # no transform for evaluation
+            "image_mask_transforms_aug": [],
+            "image_mask_transforms": []  # no transform for evaluation
         }, cache_labels=False)
 
         self.cityscapes_dataset_rescaled50 = CityscapesDatasetInterface(dataset_params={
             "batch_size": 3,
             "val_batch_size": 3,
-            "image_mask_transforms_aug": transforms.Compose([]),
-            "image_mask_transforms": transforms.Compose([Rescale(scale_factor=0.5)])  # no transform for evaluation
+            "image_mask_transforms_aug": [],
+            "image_mask_transforms": [Rescale(scale_factor=0.5)]  # no transform for evaluation
         }, cache_labels=False)
 
         self.cityscapes_dataset_rescaled75 = CityscapesDatasetInterface(dataset_params={
             "batch_size": 3,
             "val_batch_size": 3,
-            "image_mask_transforms_aug": transforms.Compose([]),
-            "image_mask_transforms": transforms.Compose([Rescale(scale_factor=0.75)])  # no transform for evaluation
+            "image_mask_transforms_aug": [],
+            "image_mask_transforms": [Rescale(scale_factor=0.75)]  # no transform for evaluation
         }, cache_labels=False)
 
         self.transfer_segmentation_dataset = SegmentationTestDatasetInterface(image_size=1024)
@@ -788,6 +793,50 @@ class PretrainedModelsTest(unittest.TestCase):
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["stdc2_seg75"], delta=0.001)
 
+    def test_pretrained_pplite_t_seg50_cityscapes(self):
+        trainer = Trainer('cityscapes_pretrained_pplite_t_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("pp_lite_t_seg50", arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"],
+                            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()
+        self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["pp_lite_t_seg50"], delta=0.001)
+
+    def test_pretrained_pplite_t_seg75_cityscapes(self):
+        trainer = Trainer('cityscapes_pretrained_pplite_t_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("pp_lite_t_seg75", arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"],
+                            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()
+        self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["pp_lite_t_seg75"], delta=0.001)
+
+    def test_pretrained_pplite_b_seg50_cityscapes(self):
+        trainer = Trainer('cityscapes_pretrained_pplite_b_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("pp_lite_b_seg50", arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"],
+                            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()
+        self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["pp_lite_b_seg50"], delta=0.001)
+
+    def test_pretrained_pplite_b_seg75_cityscapes(self):
+        trainer = Trainer('cityscapes_pretrained_pplite_b_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("pp_lite_b_seg75", arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"],
+                            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()
+        self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["pp_lite_b_seg75"], delta=0.001)
+
     def test_transfer_learning_stdc2_seg75_cityscapes(self):
         trainer = Trainer('cityscapes_pretrained_stdc2_seg75_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)