@@ -61,7 +61,7 @@ class TestDetectionDatasetCaching(unittest.TestCase):
 
         self.assertEqual(1, self._count_cached_array())
         for first_dataset, second_dataset in zip(datasets[:-1], datasets[1:]):
-            self.assertTrue(np.array_equal(first_dataset.cached_imgs, second_dataset.cached_imgs))
+            self.assertTrue(np.array_equal(first_dataset.cached_imgs_padded, second_dataset.cached_imgs_padded))
 
         self._empty_cache()
 
@@ -76,7 +76,7 @@ class TestDetectionDatasetCaching(unittest.TestCase):
 
         self.assertEqual(5, self._count_cached_array())
         for first_dataset, second_dataset in zip(datasets[:-1], datasets[1:]):
-            self.assertFalse(np.array_equal(first_dataset.cached_imgs, second_dataset.cached_imgs))
+            self.assertFalse(np.array_equal(first_dataset.cached_imgs_padded, second_dataset.cached_imgs_padded))
 
         self._empty_cache()
 

@@ -30,7 +30,7 @@ class DummyDetectionDataset(DetectionDataset):
             - Annotation 0 is made of: 3 targets of class 0, 0 of class_1 and 0 of class_2
             - Annotation 1 is made of: 2 targets of class_0, 1 of class_1 and 0 of class_2
         """
-        return {"img_path": "", "target": self.dummy_targets[sample_id]}
+        return {"img_path": "", "resized_img_shape": None, "target": self.dummy_targets[sample_id]}
 
     # DetectionDatasetV2 will call _load_image but since we don't have any image we patch this method with
     # tensor of image shape

@@ -21,7 +21,7 @@ class DummyDetectionDataset(DetectionDataset):
 
     def _load_annotation(self, sample_id: int) -> dict:
         """Load dummy annotation"""
-        return {"img_path": "", "target": torch.zeros(10, 6)}
+        return {"img_path": "", "resized_img_shape": None, "target": torch.zeros(10, 6)}
 
     # DetectionDatasetV2 will call _load_image but since we don't have any image we patch this method with
     # tensor of image shape