@@ -231,6 +231,7 @@ class MultiscalePrePredictionCallback(AbstractPrePredictionCallback):
         self.rank = None
         self.is_distributed = None
         self.sampled_imres_once = False
+        self.new_input_size = None
 
     def __call__(self, inputs, targets, batch_idx):
         if self.rank is None:
@@ -239,9 +240,9 @@ class MultiscalePrePredictionCallback(AbstractPrePredictionCallback):
             self.is_distributed = get_world_size() > 1
 
         # GENERATE A NEW SIZE AND BROADCAST IT TO THE THE OTHER RANKS SO THEY HAVE THE SAME SCALE
+        input_size = inputs.shape[2:]
         if batch_idx % self.frequency == 0:
-            tensor = torch.LongTensor(2).cuda()
-            input_size = inputs.shape[2:]
+            tensor = torch.LongTensor(2).to(inputs.device)
 
             if self.rank == 0:
                 size_factor = input_size[1] * 1.0 / input_size[0]
@@ -262,12 +263,12 @@ class MultiscalePrePredictionCallback(AbstractPrePredictionCallback):
                 dist.barrier()
                 dist.broadcast(tensor, 0)
 
-            new_input_size = (tensor[0].item(), tensor[1].item())
+            self.new_input_size = (tensor[0].item(), tensor[1].item())
 
-            scale_y = new_input_size[0] / input_size[0]
-            scale_x = new_input_size[1] / input_size[1]
-            if scale_x != 1 or scale_y != 1:
-                inputs = torch.nn.functional.interpolate(inputs, size=new_input_size, mode="bilinear", align_corners=False)
+        scale_y = self.new_input_size[0] / input_size[0]
+        scale_x = self.new_input_size[1] / input_size[1]
+        if scale_x != 1 or scale_y != 1:
+            inputs = torch.nn.functional.interpolate(inputs, size=self.new_input_size, mode="bilinear", align_corners=False)
         return inputs, targets
 
 

@@ -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()