Fix issue with nested model output as input

keras/src/models/functional_test.py

@@ -714,3 +714,19 @@ def test_list_input_with_dict_build(self):
             "The structure of `inputs` doesn't match the expected structure",
         ):
             model([x1, x2])
+
+    def test_nested_model_call_as_arg(self):
+        model_1 = Sequential([
+            Input(shape=(6,)),
+            layers.Dense(3, activation="sigmoid"),
+        ])
+
+        model_2 = Sequential([
+            Input(shape=(3,)),
+            layers.Dense(1, activation="sigmoid"),
+        ], )
+
+        combined = Model(Input(shape=(6,)), model_2(model_1(Input(shape=(6,)))), name='nested_model')
+        combined.compile(loss='binary_crossentropy', optimizer='adam')
+        output = combined.train_on_batch(np.random.normal(0, 1, (8, 6)), np.random.normal(0, 1, (8, 1)))
+        self.assertNotIsInstance(output, type(None))

keras/src/ops/function.py

@@ -160,8 +160,14 @@ def _run_through_graph(self, inputs, operation_fn, call_fn=None):
                 if not node.operation or node.is_input:
                     continue  # Input tensors already exist.
 
-                if any(id(x) not in tensor_dict for x in node.input_tensors):
-                    continue  # Node is not computable, try skipping.
+                for x in node.input_tensors:
+                    if id(x) not in tensor_dict:
+                        if id(node.outputs[0]) == id(self.outputs[0]):
+                            tensor_dict[id(x)] = inputs[0]
+                        elif x.shape == self.inputs[0].shape:
+                            tensor_dict[id(x)] = inputs[0]
+                        else:
+                            pass  # Node is not computable, try skipping.
 
                 args, kwargs = node.arguments.fill_in(tensor_dict)
                 op = operation_fn(node.operation)