From 93dda17944b6e34380897ad3480ac2218fb7398e Mon Sep 17 00:00:00 2001
From: Aleksej Solomatin <aleksejsolomatin@gmail.com>
Date: Sun, 30 Jun 2024 22:44:03 +0300
Subject: [PATCH] fix: Added support for training a multi-input model using a
 dataset.

---
 src/TensorFlowNET.Keras/Engine/Model.Evaluate.cs | 14 +++++++++++++-
 src/TensorFlowNET.Keras/Engine/Model.Fit.cs      | 13 ++++++++++++-
 2 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/src/TensorFlowNET.Keras/Engine/Model.Evaluate.cs b/src/TensorFlowNET.Keras/Engine/Model.Evaluate.cs
index b3264429..ec99d7ef 100644
--- a/src/TensorFlowNET.Keras/Engine/Model.Evaluate.cs
+++ b/src/TensorFlowNET.Keras/Engine/Model.Evaluate.cs
@@ -112,7 +112,19 @@ public Dictionary<string, float> evaluate(IDatasetV2 x, int verbose = 1, bool is
                 Steps = data_handler.Inferredsteps
             });
 
-            return evaluate(data_handler, callbacks, is_val, test_function);
+            Func<DataHandler, OwnedIterator, Dictionary<string, float>> testFunction;
+
+            if (data_handler.DataAdapter.GetDataset().structure.Length > 2 ||
+                data_handler.DataAdapter.GetDataset().FirstInputTensorCount > 1)
+            {
+                testFunction = test_step_multi_inputs_function;
+            }
+            else
+            {
+                testFunction = test_function;
+            }
+
+            return evaluate(data_handler, callbacks, is_val, testFunction);
         }
 
         /// <summary>
diff --git a/src/TensorFlowNET.Keras/Engine/Model.Fit.cs b/src/TensorFlowNET.Keras/Engine/Model.Fit.cs
index 13a1b63b..e1303513 100644
--- a/src/TensorFlowNET.Keras/Engine/Model.Fit.cs
+++ b/src/TensorFlowNET.Keras/Engine/Model.Fit.cs
@@ -179,9 +179,20 @@ public ICallback fit(IDatasetV2 dataset,
                 StepsPerExecution = _steps_per_execution
             });
 
+            Func<DataHandler, OwnedIterator, Dictionary<string, float>> trainStepFunction;
+
+            if (data_handler.DataAdapter.GetDataset().structure.Length > 2 ||
+                data_handler.DataAdapter.GetDataset().FirstInputTensorCount > 1)
+            {
+                trainStepFunction = train_step_multi_inputs_function;
+            }
+            else
+            {
+                trainStepFunction = train_step_function;
+            }
 
             return FitInternal(data_handler, epochs, validation_step, verbose, callbacks, validation_data: validation_data,
-                    train_step_func: train_step_function);
+                    train_step_func: trainStepFunction);
         }
 
         History FitInternal(DataHandler data_handler, int epochs, int validation_step, int verbose, List<ICallback> callbackList, IDatasetV2 validation_data,