diff --git a/tests/callbacks/test_debugging.py b/tests/callbacks/test_debugging.py
index aa389a6..867b0c5 100644
--- a/tests/callbacks/test_debugging.py
+++ b/tests/callbacks/test_debugging.py
@@ -38,7 +38,7 @@ def generate_data() -> Iterator[Tuple[List[int], List[int]]]:
         ),
     )
 
-    model = tf.keras.Sequential([tf.keras.Input(shape=(1)), tf.keras.layers.Dense(1)])
+    model = tf.keras.Sequential([tf.keras.Input(shape=(1,)), tf.keras.layers.Dense(1)])
     model.compile(loss="binary_crossentropy")
     model.fit(fake_data, steps_per_epoch=1, epochs=10, callbacks=[HangDebugCallback()])
 
@@ -70,7 +70,7 @@ def generate_data() -> Iterator[Tuple[List[int], List[int]]]:
         ),
     )
 
-    model = tf.keras.Sequential([tf.keras.Input(shape=(1)), tf.keras.layers.Dense(1)])
+    model = tf.keras.Sequential([tf.keras.Input(shape=(1,)), tf.keras.layers.Dense(1)])
     model.compile(loss="binary_crossentropy")
     model.fit(
         fake_data,
@@ -114,7 +114,7 @@ def generate_data() -> Iterator[Tuple[List[int], List[int]]]:
         ),
     )
 
-    model = tf.keras.Sequential([tf.keras.Input(shape=(1)), tf.keras.layers.Dense(1)])
+    model = tf.keras.Sequential([tf.keras.Input(shape=(1,)), tf.keras.layers.Dense(1)])
     model.compile(loss="binary_crossentropy")
     model.fit(
         fake_data,