update test cases and format the code

innat · Mar 6, 2024 · 8287395 · 8287395
1 parent c8aea50
commit 8287395
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/keras_cv/models/classification/video_classifier_test.py b/keras_cv/models/classification/video_classifier_test.py
@@ -86,7 +86,7 @@ def test_saved_model(self):
             ),
             num_classes=10,
         )
-        model_output = model.predict(self.input_batch)
+        model_output = model(self.input_batch)
         save_path = os.path.join(self.get_temp_dir(), "video_classifier.keras")
         model.save(save_path)
         restored_model = keras.models.load_model(save_path)
@@ -95,7 +95,7 @@ def test_saved_model(self):
         self.assertIsInstance(restored_model, VideoClassifier)
 
         # Check that output matches.
-        restored_output = restored_model.predict(self.input_batch)
+        restored_output = restored_model(self.input_batch)
         self.assertAllClose(
             ops.convert_to_numpy(model_output),
             ops.convert_to_numpy(restored_output),