message: Add back flattening of inputs in v2 training loop.
diff: @@ -62,7 +62,7 @@ def _make_execution_function(model, mode):
  def distributed_function(input_iterator):
  """A single step of the distributed execution across replicas."""
  - x, y, sample_weights = _prepare_feed_values(input_iterator, mode)
  + x, y, sample_weights = _prepare_feed_values(model, input_iterator, mode)
  # Call `Model.{train,test,predict}_on_batch` on every replica passing
  # PerReplicas as arguments. On every replica inside this call, each
  # PerReplica object will return the value for that replica. The outputs
  @@ -92,13 +92,14 @@ def _non_none_constant_value(v):
  return constant_value if constant_value is not None else v
  -def _prepare_feed_values(inputs, mode):
  +def _prepare_feed_values(model, inputs, mode):
  """Prepare feed values to the model execution function.
  Arguments:
  - inputs: An iterator of model inputs, targets, and sample_weights. model
  - inputs may be lists, single values, or dicts mapping input feed names to
  - values.
  + model: Model to prepare feed values for.
  + inputs: An iterator of model inputs, targets, and sample_weights.
  + model inputs may be lists, single values, or dicts mapping input feed
  + names to values.
  mode: One of ModeKeys.TRAIN/ModeKeys.TEST/ModeKeys.PREDICT.
  Returns:
  @@ -109,6 +110,14 @@ def _prepare_feed_values(inputs, mode):
  """
  inputs, targets, sample_weights = _get_input_from_iterator(inputs)
  + # When the inputs are dict, then we want to flatten it in the same order as
  + # the input layers, such that the data are fed into the input layers in the
  + # correct order.
  + if isinstance(inputs, dict):
  + inputs = [inputs[key] for key in model._feed_input_names]
  + else:
  + inputs = training_utils.ModelInputs(inputs).as_list()
  +
  if mode == ModeKeys.PREDICT:
  sample_weights = []
  targets = []
  