🐛 bug fix 2

innat · Mar 22, 2024 · 5787cbb · 5787cbb
1 parent d083703
commit 5787cbb
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/test/test_classifier.py b/test/test_classifier.py
@@ -19,20 +19,22 @@ def setUp(self):
         ).batch(4)
 
     def test_valid_call(self):
+        input_batch = np.ones(shape=(2, 8, 256, 256, 3))
         model = VideoSwinT(
             input_shape=(8, 256, 256, 3),
             include_rescaling=False,
             num_classes=10,
         )
-        model(self.input_batch)
+        model(input_batch)
 
     def test_valid_call_non_square_shape(self):
+        input_batch = np.ones(shape=(2, 8, 224, 256, 3))
         model = VideoSwinT(
             input_shape=(8, 224, 256, 3),
             include_rescaling=False,
             num_classes=10,
         )
-        model(self.input_batch)
+        model(input_batch)
 
     @parameterized.named_parameters(
         ("jit_compile_false", False), ("jit_compile_true", True)
@@ -54,13 +56,14 @@ def test_classifier_fit(self, jit_compile):
 
     @parameterized.named_parameters(("avg_pooling", "avg"), ("max_pooling", "max"))
     def test_pooling_arg_call(self, pooling):
+        input_batch = np.ones(shape=(2, 8, 224, 224, 3))
         model = VideoSwinT(
             input_shape=(8, 224, 224, 3),
             include_rescaling=True,
             num_classes=10,
             pooling=pooling,
         )
-        model(self.input_batch)
+        model(input_batch)
 
     @pytest.mark.large  # Saving is slow, so mark these large.
     def test_saved_model(self):