Repository: keras-team/keras
Issue: Strange behaviour of model.predict

**System information**.
- Have I written custom code (as opposed to using a stock example script provided in Keras): yes
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Google Colab
- TensorFlow installed from (source or binary): provided
- TensorFlow version (use command below): 2.8.0
- Python version: 3.7
- Bazel version (if compiling from source):
- GPU model and memory: no
- Exact command to reproduce:

**Describe the problem**.

Prediction for models with structured output sometimes fails on tf datasets with static shapes when `steps_per_execution` is greater than one. The error message is not descriptive.

After a successful prediction, the predictions that were failing don't fail anymore.

**Describe the current behavior**.

Prediction fails with  `'outputs' has shape <...> before the loop, but shape <...> after one iteration`

**Describe the expected behavior**.

Either a consistent successful execution or a meaningful error message.

**[Contributing](https://github.com/keras-team/keras/blob/master/CONTRIBUTING.md)**.

- Do you want to contribute a PR? (yes/no): no
- If yes, please read [this page](https://github.com/keras-team/keras/blob/master/CONTRIBUTING.md) for instructions
- Briefly describe your candidate solution(if contributing):

**Standalone code to reproduce the issue**.

Example notebook: https://colab.research.google.com/gist/itmo153277/d921710f4876f3883da3b37dc7122599/testpredict.ipynb

**Source code / logs**.

See notebook above

Changed files:
keras/engine/training_test.py, keras/engine/training.py

diff --git a/keras/engine/training.py b/keras/engine/training.py
index 51c650ee79..e906a737b2 100644
--- a/keras/engine/training.py
+++ b/keras/engine/training.py
@@ -1851,9 +1851,10 @@ class Model(base_layer.Layer, version_utils.ModelVersionSelector):
         outputs = step_function(self, iterator)
         for _ in tf.range(self._steps_per_execution - 1):
           tf.autograph.experimental.set_loop_options(
-              shape_invariants=[(
-                  t, tf_utils.get_tensor_spec(t, dynamic_batch=True).shape)
-                                for t in tf.nest.flatten(outputs)])
+              shape_invariants=[(outputs, tf.nest.map_structure(
+                  lambda t: tf_utils.get_tensor_spec(
+                      t, dynamic_batch=True).shape,
+                  outputs))])
           step_outputs = step_function(self, iterator)
           outputs = tf.nest.map_structure(lambda t1, t2: concat([t1, t2]),
                                           outputs, step_outputs)
diff --git a/keras/engine/training_test.py b/keras/engine/training_test.py
index 9dab225115..228e0f9914 100644
--- a/keras/engine/training_test.py
+++ b/keras/engine/training_test.py
@@ -2078,6 +2078,31 @@ class TestExceptionsAndWarnings(test_combinations.TestCase):
                                 'Unexpected result of `predict_function`.*'):
       model.predict(np.array([]))
 
+  @test_combinations.run_all_keras_modes(always_skip_v1=True)
+  @parameterized.named_parameters(
+      ('dynamic', 0, False),
+      ('dynamic_multistep', 10, False),
+      ('static', 0, True),
+      ('static_multistep', 10, True),
+  )
+  def test_predict_structured(self, spe, static_batch):
+    inputs = layers_module.Input(shape=(2,))
+    outputs = layers_module.Dense(2)(inputs)
+    model = training_module.Model(
+        inputs=inputs,
+        outputs={'out': outputs},
+    )
+    model.compile(
+        loss='mse',
+        steps_per_execution=spe,
+        run_eagerly=test_utils.should_run_eagerly(),
+    )
+    xdata = np.random.uniform(size=(8, 2)).astype(np.float32)
+    dataset = tf.data.Dataset.from_tensor_slices((xdata, xdata))
+    dataset = dataset.batch(8, drop_remainder=static_batch)
+    ret = model.predict(dataset, steps=1)
+    tf.nest.assert_same_structure(ret, {'out': ''})
+
   @test_combinations.run_all_keras_modes(always_skip_v1=True)
   def test_on_batch_error_inconsistent_batch_size(self):
     input_node1 = layers_module.Input(shape=(5,))