diff --git a/sdks/python/apache_beam/ml/inference/base_test.py b/sdks/python/apache_beam/ml/inference/base_test.py
index 64fd73682e13..66e85ce163e7 100644
--- a/sdks/python/apache_beam/ml/inference/base_test.py
+++ b/sdks/python/apache_beam/ml/inference/base_test.py
@@ -1141,7 +1141,7 @@ def test_run_inference_with_iterable_side_input(self):
             accumulation_mode=trigger.AccumulationMode.DISCARDING))
 
     test_pipeline.options.view_as(StandardOptions).streaming = True
-    with self.assertRaises(ValueError) as e:
+    with self.assertRaises(Exception) as e:
       _ = (
           test_pipeline
           | beam.Create([1, 2, 3, 4])
@@ -1165,7 +1165,7 @@ def test_run_inference_with_iterable_side_input_multi_process_shared(self):
             accumulation_mode=trigger.AccumulationMode.DISCARDING))
 
     test_pipeline.options.view_as(StandardOptions).streaming = True
-    with self.assertRaises(ValueError) as e:
+    with self.assertRaises(Exception) as e:
       _ = (
           test_pipeline
           | beam.Create([1, 2, 3, 4])