Update test_debugging.py

spotify · Aug 8, 2024 · 5c1b72f · 5c1b72f
1 parent 26e62b3
commit 5c1b72f
Showing 1 changed file with 3 additions and 3 deletions.
diff --git 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,