diff --git a/sdks/python/apache_beam/ml/transforms/base_test.py b/sdks/python/apache_beam/ml/transforms/base_test.py
index 64d0afe955d0..c7b904c09444 100644
--- a/sdks/python/apache_beam/ml/transforms/base_test.py
+++ b/sdks/python/apache_beam/ml/transforms/base_test.py
@@ -150,8 +150,8 @@ def test_ml_transform_on_list_dict(self):
       ),
       param(
           input_data=[{
-              'x': np.array([1], dtype=np.int64),
-              'y': np.array([2.0], dtype=np.float32),
+              'x': np.int64(1),
+              'y': np.float32(2.0),
           }],
           input_types={
               'x': np.int32, 'y': np.float32