update test method to swin backbone

innat · Mar 24, 2024 · 251495b · 251495b
1 parent 204e4b1
commit 251495b
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/keras_cv/models/backbones/video_swin/video_swin_backbone_test.py b/keras_cv/models/backbones/video_swin/video_swin_backbone_test.py
@@ -82,3 +82,13 @@ def test_can_run_on_gray_video(self):
         y = np.zeros((1, 48, 3, 3, 768))
         model.compile(optimizer="adam", loss="mse", metrics=["mse"])
         model.fit(x, y, epochs=1)
+
+    @pytest.mark.extra_large
+    def test_can_run_non_square_shape(self):
+        input_batch = np.ones(shape=(2, 8, 224, 256, 3))
+        model = VideoSwinBackbone(
+            input_shape=(8, 224, 256, 3),
+            include_rescaling=False,
+            num_classes=10,
+        )
+        model(input_batch)