fix shape in tests

keras-team · Oct 20, 2024 · c7a6166 · c7a6166
1 parent 225942d
commit c7a6166
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/keras_hub/src/models/mit/mit_backbone_test.py b/keras_hub/src/models/mit/mit_backbone_test.py
@@ -9,7 +9,7 @@ class MiTBackboneTest(TestCase):
     def setUp(self):
         self.init_kwargs = {
             "depths": [2, 2],
-            "image_shape": (16, 16, 3),
+            "image_shape": (32, 32, 3),
             "hidden_dims": [4, 8],
             "num_layers": 2,
             "blockwise_num_heads": [1, 2],
@@ -18,7 +18,7 @@ def setUp(self):
             "patch_sizes": [7, 3],
             "strides": [4, 2],
         }
-        self.input_size = 16
+        self.input_size = 32
         self.input_data = np.ones(
             (2, self.input_size, self.input_size, 3), dtype="float32"
         )
@@ -28,9 +28,9 @@ def test_backbone_basics(self):
             cls=MiTBackbone,
             init_kwargs=self.init_kwargs,
             input_data=self.input_data,
-            expected_output_shape=(2, 2, 2, 8),
+            expected_output_shape=(2, 4, 4, 8),
             expected_pyramid_output_keys=["P1", "P2"],
-            expected_pyramid_image_sizes=[(4, 4), (2, 2)],
+            expected_pyramid_image_sizes=[(8, 8), (4, 4)],
             run_quantization_check=False,
             run_mixed_precision_check=False,
             run_data_format_check=False,