@@ -26,6 +26,7 @@ from tests.unit_tests.forward_pass_prep_fn_test import ForwardpassPrepFNTest
 from tests.unit_tests.mask_loss_test import MaskAttentionLossTest
 from tests.unit_tests.detection_sub_sampling_test import TestDetectionDatasetSubsampling
 from tests.unit_tests.detection_sub_classing_test import TestDetectionDatasetSubclassing
+from tests.unit_tests.multi_scaling_test import MultiScaleTest
 
 
 class CoreUnitTestSuiteRunner:
@@ -74,6 +75,7 @@ class CoreUnitTestSuiteRunner:
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(IoULossTest))
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(TestDetectionDatasetSubsampling))
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(TestDetectionDatasetSubclassing))
+        self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(MultiScaleTest))
 
     def _add_modules_to_end_to_end_tests_suite(self):
         """

import unittest

import torch
from super_gradients.training.datasets.datasets_utils import DetectionMultiscalePrePredictionCallback


class MultiScaleTest(unittest.TestCase):

    def setUp(self) -> None:
        self.size = (1024, 512)
        self.batch_size = 12
        self.change_frequency = 10
        self.multiscale_callback = DetectionMultiscalePrePredictionCallback(change_frequency=self.change_frequency)

    def _create_batch(self):
        inputs = torch.rand((self.batch_size, 3, self.size[0], self.size[1])) * 255
        targets = torch.cat([torch.tensor([[[0, 0, 10, 10, 0]]]) for _ in range(self.batch_size)], 0)
        return inputs, targets

    def test_multiscale_keep_state(self):
        """Check that the multiscale keeps in memory the new size to use between the size swaps"""

        for i in range(5):
            post_multiscale_input_shapes = []
            for j in range(self.change_frequency):
                inputs, targets = self._create_batch()
                post_multiscale_input, _ = self.multiscale_callback(inputs, targets, batch_idx=i * self.change_frequency + j)
                post_multiscale_input_shapes.append(list(post_multiscale_input.shape))

                # The shape should be the same for a given between k * self.change_frequency and (k+1)*self.change_frequency
                self.assertListEqual(post_multiscale_input_shapes[0], post_multiscale_input_shapes[-1])


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